add retrieve by poc node in the graph

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

@@ -200,56 +200,52 @@ Graphs and their HTML embedding:
 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. 
+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.
 
-
-La table des matières est structurée de cette façon : 
+The table of contents is structured like that : 
 {{
   "level": 0,
-  "Chapitre 1": {{}},
-  "Chapitre 2" : {{
+  "Chapter 1": {{}},
+  "Chapter 2" : {{
     "level": 1,
-    "Chapitre 2.1": {{
+    "Chapter 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.
+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 ###
-- 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**.
+- 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.
 
 --------------------
-Question de l'utilisateur : 
+User question : 
 {query}
 
-Liste des documents avec leurs tables des matières : 
+List of documents with their table of contents : 
 {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
+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 : Assure-toi que les niveaux dans la réponse sont exactement ceux de la table des matières.**
+**IMPORTANT : Make sure that the levels of the answer are exactly the same as the ones in the table of contents**
 
-Exemple de réponse JSON :
+Example of a JSON response:
 [
   {{
     "document": "Document A",
-    "chapter": "Chapitre 1",
+    "chapter": "Chapter 1",
   }},
   {{
     "document": "Document B",
-    "chapter": "Chapitre 1.1",
+    "chapter": "Chapter 5",
   }}
 ]
-"""
+"""

climateqa/engine/chains/retrieve_documents.py

@@ -127,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,
@@ -177,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(
@@ -241,7 +336,6 @@ async def get_IPCC_relevant_documents(
             "report_type": { "$nin":["SPM"]},
         }
         docs_full = vectorstore.similarity_search_with_score(query=query,filter = filters_full,k = k_total)
-        docs_full = vectorstore.similarity_search_with_score(query=query,filter = filters_full,k = k_total)
         
         if search_figures:
             # Images
@@ -280,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( 
@@ -344,13 +439,10 @@ async def retrieve_documents(
             reports = reports,
         )
     
-    # 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(
+            docs_question_dict = await get_POC_documents_by_ToC_relevant_documents(
                 query=question,
                 tocs = tocs,
                 vectorstore=vectorstore,
@@ -469,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):
@@ -536,124 +654,6 @@ 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,

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)