Merged in feature/improve_parsing_and_retrieval (pull request #3)

app.py

@@ -66,7 +66,7 @@ user_id = create_user_id()
 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"))
+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":
@@ -75,7 +75,7 @@ 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)#TODO put back default 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
 
 
 async def chat(query, history, audience, sources, reports, relevant_content_sources_selection, search_only):

climateqa/engine/chains/prompts.py

@@ -198,4 +198,54 @@ Graphs and their HTML embedding:
 {format_instructions}
 
 Output the result as json with a key "graphs" containing a list of dictionaries of the relevant graphs with keys 'embedding', 'category', and 'source'. Do not modify the graph HTML embedding, the category or the source. Do not put any message or text before or after the JSON output.
-"""
+"""
+
+retrieve_chapter_prompt_template = """Given the user question and a list of documents with their table of contents, retrieve the 5 most relevant level 0 chapters which could help to answer to the question while taking account their sub-chapters.
+
+The table of contents is structured like that : 
+{{
+  "level": 0,
+  "Chapter 1": {{}},
+  "Chapter 2" : {{
+    "level": 1,
+    "Chapter 2.1": {{
+      ...
+    }}
+  }},
+}}
+
+Here level is the level of the chapter. For example, Chapter 1 and Chapter 2 are at level 0, and Chapter 2.1 is at level 1.
+
+### Guidelines ###
+- Keep all the list of documents that is given to you
+- Each chapter must keep **EXACTLY** its assigned level in the table of contents. **DO NOT MODIFY THE LEVELS. ** 
+- Check systematically the level of a chapter before including it in the answer.
+- Return **valid JSON** result.
+
+--------------------
+User question : 
+{query}
+
+List of documents with their table of contents : 
+{doc_list}
+
+--------------------
+
+Return a JSON result with a list of relevant chapters with the following keys **WITHOUT** the json markdown indicator ```json at the beginning:
+- "document" : the document in which we can find the chapter
+- "chapter" : the title of the chapter
+
+**IMPORTANT : Make sure that the levels of the answer are exactly the same as the ones in the table of contents**
+
+Example of a JSON response:
+[
+  {{
+    "document": "Document A",
+    "chapter": "Chapter 1",
+  }},
+  {{
+    "document": "Document B",
+    "chapter": "Chapter 5",
+  }}
+]
+"""

climateqa/engine/chains/retrieve_documents.py

@@ -15,6 +15,14 @@ from ..utils import log_event
 from langchain_core.vectorstores import VectorStore
 from typing import List
 from langchain_core.documents.base import Document
+from ..llm import get_llm
+from .prompts import retrieve_chapter_prompt_template
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from ..vectorstore import get_pinecone_vectorstore
+from ..embeddings import get_embeddings_function
+
+
 import asyncio
 
 from typing import Any, Dict, List, Tuple
@@ -119,6 +127,21 @@ def remove_duplicates_chunks(docs):
             result.append(doc)
     return result
 
+def get_ToCs(version: str) :
+
+    filters_text = {
+        "chunk_type":"toc",
+        "version": version
+    }
+    embeddings_function = get_embeddings_function()
+    vectorstore = get_pinecone_vectorstore(embeddings_function, index_name="climateqa-v2")
+    tocs = vectorstore.similarity_search_with_score(query="",filter = filters_text)
+
+    # remove duplicates or almost duplicates
+    tocs = remove_duplicates_chunks(tocs)
+    
+    return tocs
+
 async def get_POC_relevant_documents(
     query: str,
     vectorstore:VectorStore,
@@ -169,6 +192,86 @@ async def get_POC_relevant_documents(
         "docs_question" : docs_question,
         "docs_images" : docs_images
     }
+
+async def get_POC_documents_by_ToC_relevant_documents(
+    query: str,
+    tocs: list,
+    vectorstore:VectorStore,
+    version: str,
+    sources:list = ["Acclimaterra","PCAET","Plan Biodiversite"],
+    search_figures:bool = False,
+    search_only:bool = False,
+    k_documents:int = 10,
+    threshold:float = 0.6,
+    k_images: int = 5,
+    reports:list = [],
+    min_size:int = 200,
+    proportion: float = 0.5,
+) :
+    """
+        Args:
+            - tocs : list with the table of contents of each document
+            - version : version of the parsed documents (e.g. "v4")
+            - proportion : share of documents retrieved using ToCs
+    """
+    # Prepare base search kwargs
+    filters = {}
+    docs_question = []
+    docs_images = []
+
+    # TODO add source selection
+    # if len(reports) > 0:
+    #     filters["short_name"] = {"$in":reports}
+    # else:
+    #     filters["source"] = { "$in": sources}
+    
+    k_documents_toc = round(k_documents * proportion)
+    
+    relevant_tocs = await get_relevant_toc_level_for_query(query, tocs)
+    
+    print(f"Relevant ToCs : {relevant_tocs}")
+    # Transform the ToC dict {"document": str, "chapter": str} into a list of string 
+    toc_filters = [toc['chapter'] for toc in relevant_tocs]
+
+    filters_text_toc = {
+        **filters,
+        "chunk_type":"text",
+        "toc_level0": {"$in": toc_filters},
+        "version": version
+        # "report_type": {}, # TODO  to be completed to choose the right documents / chapters according to the analysis of the question
+    }
+    
+    docs_question = vectorstore.similarity_search_with_score(query=query,filter = filters_text_toc,k = k_documents_toc)
+
+    filters_text = {
+        **filters,
+        "chunk_type":"text",
+        "version": version
+        # "report_type": {}, # TODO  to be completed to choose the right documents / chapters according to the analysis of the question
+    }
+
+    docs_question += vectorstore.similarity_search_with_score(query=query,filter = filters_text,k = k_documents - k_documents_toc)
+
+    # remove duplicates or almost duplicates
+    docs_question = remove_duplicates_chunks(docs_question)
+    docs_question = [x for x in docs_question if x[1] > threshold]
+    
+    if search_figures:
+        # Images
+        filters_image = {
+            **filters,
+            "chunk_type":"image"
+        }
+        docs_images = vectorstore.similarity_search_with_score(query=query,filter = filters_image,k = k_images)
+        
+    docs_question, docs_images = _add_metadata_and_score(docs_question), _add_metadata_and_score(docs_images)
+    
+    docs_question = [x for x in docs_question if len(x.page_content) > min_size]
+    
+    return {
+        "docs_question" : docs_question,
+        "docs_images" : docs_images
+    }
     
 
 async def get_IPCC_relevant_documents(
@@ -271,6 +374,7 @@ def concatenate_documents(index, source_type, docs_question_dict, k_by_question,
     return docs_question, images_question
 
 
+
 # The chain callback is not necessary, but it propagates the langchain callbacks to the astream_events logger to display intermediate results
 # @chain
 async def retrieve_documents( 
@@ -279,6 +383,7 @@ async def retrieve_documents(
     source_type: str, 
     vectorstore: VectorStore,
     reranker: Any,
+    version: str = "",
     search_figures: bool = False,
     search_only: bool = False,
     reports: list = [],
@@ -286,7 +391,9 @@ async def retrieve_documents(
     k_images_by_question: int = 5,
     k_before_reranking: int = 100,
     k_by_question: int = 5,
-    k_summary_by_question: int = 3
+    k_summary_by_question: int = 3,
+    tocs: list = [],
+    by_toc=False
 ) -> Tuple[List[Document], List[Document]]:
     """
     Unpack the first question of the remaining questions, and retrieve and rerank corresponding documents, based on the question and selected_sources 
@@ -316,6 +423,7 @@ async def retrieve_documents(
 
     print(f"""---- Retrieve documents from {current_question["source_type"]}----""")
 
+    
     if source_type == "IPx":
         docs_question_dict = await get_IPCC_relevant_documents(
             query  = question,
@@ -331,19 +439,36 @@ async def retrieve_documents(
             reports = reports,
         )
     
-    if source_type == "POC":
-        docs_question_dict = await get_POC_relevant_documents(
-            query = question,
-            vectorstore=vectorstore,
-            search_figures = search_figures,
-            sources = sources,
-            threshold = 0.5,
-            search_only = search_only,
-            reports = reports,
-            min_size= 200,
-            k_documents= k_before_reranking,
-            k_images= k_by_question
-        )
+    if source_type == 'POC':
+        if by_toc == True:
+            print("---- Retrieve documents by ToC----")
+            docs_question_dict = await get_POC_documents_by_ToC_relevant_documents(
+                query=question,
+                tocs = tocs,
+                vectorstore=vectorstore,
+                version=version,
+                search_figures = search_figures,
+                sources = sources,
+                threshold = 0.5,
+                search_only = search_only,
+                reports = reports,
+                min_size= 200,
+                k_documents= k_before_reranking,
+                k_images= k_by_question
+            )
+        else : 
+            docs_question_dict = await get_POC_relevant_documents(
+                query = question,
+                vectorstore=vectorstore,
+                search_figures = search_figures,
+                sources = sources,
+                threshold = 0.5,
+                search_only = search_only,
+                reports = reports,
+                min_size= 200,
+                k_documents= k_before_reranking,
+                k_images= k_by_question
+            )
 
     # Rerank
     if reranker is not None and rerank_by_question:
@@ -382,7 +507,10 @@ async def retrieve_documents_for_all_questions(
     reranker, 
     rerank_by_question=True, 
     k_final=15, 
-    k_before_reranking=100
+    k_before_reranking=100,
+    version: str = "",
+    tocs: list[dict] = [],
+    by_toc: bool = False
 ):
     """
     Retrieve documents in parallel for all questions.
@@ -403,6 +531,7 @@ async def retrieve_documents_for_all_questions(
     k_images_by_question = _get_k_images_by_question(n_questions)
     k_before_reranking=100
     
+    print(f"Source type here is {source_type}")
     tasks = [
         retrieve_documents(
             current_question=question,
@@ -417,7 +546,10 @@ async def retrieve_documents_for_all_questions(
             k_images_by_question=k_images_by_question,
             k_before_reranking=k_before_reranking,
             k_by_question=k_by_question,
-            k_summary_by_question=k_summary_by_question
+            k_summary_by_question=k_summary_by_question,
+            tocs=tocs,
+            version=version,
+            by_toc=by_toc
         )
         for i, question in enumerate(questions_list) if i in to_handle_questions_index
     ]
@@ -429,6 +561,32 @@ async def retrieve_documents_for_all_questions(
         new_state["related_contents"].extend(images_question)
     return new_state
 
+# ToC Retriever
+async def get_relevant_toc_level_for_query(
+    query: str,
+    tocs: list[Document],
+) -> list[dict] : 
+
+    doc_list = []
+    for doc in tocs:
+        doc_name = doc[0].metadata['name']
+        toc = doc[0].page_content
+        doc_list.append({'document': doc_name, 'toc': toc})
+
+    llm = get_llm(provider="openai",max_tokens = 1024,temperature = 0.0)
+
+    prompt = ChatPromptTemplate.from_template(retrieve_chapter_prompt_template)
+    chain = prompt | llm | StrOutputParser()
+    response = chain.invoke({"query": query, "doc_list": doc_list})
+
+    try: 
+        relevant_tocs = eval(response)
+    except Exception as e:
+        print(f" Failed to parse the result because of : {e}")
+
+    return relevant_tocs
+
+
 def make_IPx_retriever_node(vectorstore,reranker,llm,rerank_by_question=True, k_final=15, k_before_reranking=100, k_summary=5):
     
     async def retrieve_IPx_docs(state, config):
@@ -496,4 +654,56 @@ def make_POC_retriever_node(vectorstore,reranker,llm,rerank_by_question=True, k_
     return retrieve_POC_docs_node
 
 
+def make_POC_by_ToC_retriever_node(
+        vectorstore: VectorStore, 
+        reranker,
+        llm, 
+        version: str = "", 
+        rerank_by_question=True, 
+        k_final=15, 
+        k_before_reranking=100, 
+        k_summary=5,
+    ):
+    
+    async def retrieve_POC_docs_node(state, config):
+        if "POC region" not in state["relevant_content_sources_selection"]  :  
+            return {}
+        
+        search_figures = "Figures (IPCC/IPBES)" in state["relevant_content_sources_selection"]
+        search_only = state["search_only"]
+        search_only = state["search_only"]
+        reports = state["reports"]
+        questions_list = state["questions_list"]
+        n_questions=state["n_questions"]["total"]
+
+        tocs = get_ToCs(version=version)
+
+        source_type = "POC"
+        POC_questions_index = [i for i, x in enumerate(state["questions_list"]) if x["source_type"] == "POC"]
+        
+        state = await retrieve_documents_for_all_questions(
+            search_figures=search_figures,
+            search_only=search_only,
+            config=config,
+            reports=reports,
+            questions_list=questions_list,
+            n_questions=n_questions,
+            source_type=source_type,
+            to_handle_questions_index=POC_questions_index,
+            vectorstore=vectorstore,
+            reranker=reranker,
+            rerank_by_question=rerank_by_question,
+            k_final=k_final,
+            k_before_reranking=k_before_reranking,
+            tocs=tocs,
+            version=version,
+            by_toc=True
+        )
+        return state
+    
+    return retrieve_POC_docs_node
+
+
+
+
 

climateqa/engine/graph.py

@@ -19,7 +19,7 @@ from .chains.answer_ai_impact import make_ai_impact_node
 from .chains.query_transformation import make_query_transform_node
 from .chains.translation import make_translation_node
 from .chains.intent_categorization import make_intent_categorization_node
-from .chains.retrieve_documents import make_IPx_retriever_node, make_POC_retriever_node
+from .chains.retrieve_documents import make_IPx_retriever_node, make_POC_retriever_node, make_POC_by_ToC_retriever_node
 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
@@ -211,8 +211,23 @@ def make_graph_agent(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_regi
     app = workflow.compile()
     return app
 
-def make_graph_agent_poc(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_region, reranker, threshold_docs=0.2):
+def make_graph_agent_poc(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_region, reranker, version:str, threshold_docs=0.2):
+    """_summary_
+
+    Args:
+        llm (_type_): _description_
+        vectorstore_ipcc (_type_): _description_
+        vectorstore_graphs (_type_): _description_
+        vectorstore_region (_type_): _description_
+        reranker (_type_): _description_
+        version (str): version of the parsed documents (e.g "v4")
+        threshold_docs (float, optional): _description_. Defaults to 0.2.
+
+    Returns:
+        _type_: _description_
+    """
     
+
     workflow = StateGraph(GraphState)
 
     # Define the node functions
@@ -223,7 +238,8 @@ def make_graph_agent_poc(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_
     answer_ai_impact = make_ai_impact_node(llm)
     retrieve_documents = make_IPx_retriever_node(vectorstore_ipcc, reranker, llm)
     retrieve_graphs = make_graph_retriever_node(vectorstore_graphs, reranker)
-    retrieve_local_data = make_POC_retriever_node(vectorstore_region, reranker, llm)
+    # retrieve_local_data = make_POC_retriever_node(vectorstore_region, reranker, llm)
+    retrieve_local_data = make_POC_by_ToC_retriever_node(vectorstore_region, reranker, llm, version=version) 
     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)