retrieve by ToC

climateqa/engine/chains/prompts.py

@@ -197,4 +197,58 @@ 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 = """A partir de la question de l'utilisateur et d'une liste de documents avec leurs tables des matières, retrouve les 5 chapitres de niveau 0 les plus pertinents qui pourraient aider à répondre à la question tout en prenant bien en compte leurs sous-chapitres. 
+
+
+La table des matières est structurée de cette façon : 
+{{
+  "level": 0,
+  "Chapitre 1": {{}},
+  "Chapitre 2" : {{
+    "level": 1,
+    "Chapitre 2.1": {{
+      ...
+    }}
+  }},
+}}
+
+Ici level correspond au niveau du chapitre. Ici par exemple Chapitre 1 et Chapitre 2 sont au niveau 0, et Chapitre 2.1 est au niveau 1
+
+-----------------------
+Suis impérativement cette guideline ci-dessous.
+
+### Guidelines ###
+- Retiens bien la liste complète des documents.
+- Chaque chapitre doit conserver **EXACTEMENT** le niveau qui lui est attribué dans la table des matières. **NE MODIFIE PAS LES NIVEAUX.**
+- Vérifie systématiquement le niveau d’un chapitre avant de l’inclure dans la réponse.
+- Retourne un résultat **JSON valide**.
+
+--------------------
+Question de l'utilisateur : 
+{query}
+
+Liste des documents avec leurs tables des matières : 
+{doc_list}
+
+--------------------
+
+Retourne le résultat en JSON contenant une liste des chapitres pertinents avec les clés suivantes sans l'indicateur markdown du json:
+- "document" : le document contenant le chapitre
+- "chapter" : le titre du chapitre
+
+**IMPORTANT : Assure-toi que les niveaux dans la réponse sont exactement ceux de la table des matières.**
+
+Exemple de réponse JSON :
+[
+  {{
+    "document": "Document A",
+    "chapter": "Chapitre 1",
+  }},
+  {{
+    "document": "Document B",
+    "chapter": "Chapitre 1.1",
+  }}
+]
 """

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
@@ -280,6 +288,7 @@ async def retrieve_documents(
     source_type: str, 
     vectorstore: VectorStore,
     reranker: Any,
+    version: str = "",
     search_figures: bool = False,
     search_only: bool = False,
     reports: list = [],
@@ -287,7 +296,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 
@@ -317,6 +328,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,
@@ -332,19 +344,39 @@ 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 source_type == 'POC':
+        if by_toc == True:
+            print("---- Retrieve documents by ToC----")
+            docs_question_dict = await get_POC_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:
@@ -370,7 +402,20 @@ async def retrieve_documents(
     return docs_question, images_question
     
 
-async def retrieve_documents_for_all_questions(state, config, source_type, to_handle_questions_index, vectorstore, reranker, rerank_by_question=True, k_final=15, k_before_reranking=100):
+async def retrieve_documents_for_all_questions(
+        state, 
+        config, 
+        source_type, 
+        to_handle_questions_index, 
+        vectorstore: VectorStore, 
+        reranker,
+        version: str = "",
+        rerank_by_question: bool=True, 
+        k_final: int=15, 
+        k_before_reranking:int =100, 
+        tocs=[],
+        by_toc=False
+    ):
     """
     Retrieve documents in parallel for all questions.
     """
@@ -388,6 +433,7 @@ async def retrieve_documents_for_all_questions(state, config, source_type, to_ha
     k_images_by_question = _get_k_images_by_question(state["n_questions"]["total"])
     k_before_reranking=100
     
+    print(f"Source type here is {source_type}")
     tasks = [
         retrieve_documents(
             current_question=question,
@@ -402,7 +448,10 @@ async def retrieve_documents_for_all_questions(state, config, source_type, to_ha
             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(state["questions_list"]) if i in to_handle_questions_index
     ]
@@ -463,4 +512,164 @@ def make_POC_retriever_node(vectorstore,reranker,llm,rerank_by_question=True, k_
     return retrieve_POC_docs_node
 
 
+# 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
+
+
+async def get_POC_by_ToC_relevant_documents(
+    query: str,
+    tocs: list[str],
+    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,
+) :
+    # 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
+    }
+    
+
+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
+
+
+
+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 {}
+        
+
+        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(
+            state=state,
+            config=config,
+            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
+
+
+
+