climateqa/engine/chains/retriever.py

import sys
import os
from contextlib import contextmanager

from ..reranker import rerank_docs
from ...knowledge.retriever import ClimateQARetriever




def divide_into_parts(target, parts):
    # Base value for each part
    base = target // parts
    # Remainder to distribute
    remainder = target % parts
    # List to hold the result
    result = []
    
    for i in range(parts):
        if i < remainder:
            # These parts get base value + 1
            result.append(base + 1)
        else:
            # The rest get the base value
            result.append(base)
    
    return result


@contextmanager
def suppress_output():
    # Open a null device
    with open(os.devnull, 'w') as devnull:
        # Store the original stdout and stderr
        old_stdout = sys.stdout
        old_stderr = sys.stderr
        # Redirect stdout and stderr to the null device
        sys.stdout = devnull
        sys.stderr = devnull
        try:
            yield
        finally:
            # Restore stdout and stderr
            sys.stdout = old_stdout
            sys.stderr = old_stderr



def make_retriever_node(vectorstore,reranker,rerank_by_question=True, k_final=15, k_before_reranking=100, k_summary=5):

    def retrieve_documents(state):
        
        POSSIBLE_SOURCES = ["IPCC","IPBES","IPOS"] # ,"OpenAlex"]
        questions = state["questions"]
        
        # Use sources from the user input or from the LLM detection
        if "sources_input" not in state or state["sources_input"] is None:
            sources_input = ["auto"]
        else:
            sources_input = state["sources_input"]
        auto_mode = "auto" in sources_input

        # There are several options to get the final top k
        # Option 1 - Get 100 documents by question and rerank by question
        # Option 2 - Get 100/n documents by question and rerank the total
        if rerank_by_question:
            k_by_question = divide_into_parts(k_final,len(questions))
        
        docs = []
        
        for i,q in enumerate(questions):
            
            sources = q["sources"]
            question = q["question"]
            
            # If auto mode, we use the sources detected by the LLM
            if auto_mode:
                sources = [x for x in sources if x in POSSIBLE_SOURCES]
                
            # Otherwise, we use the config
            else:
                sources = sources_input
                
            # Search the document store using the retriever
            # Configure high top k for further reranking step
            retriever = ClimateQARetriever(
                vectorstore=vectorstore,
                sources = sources,
                # reports = ias_reports,
            min_size = 200,
            k_summary = k_summary,k_total = k_before_reranking,
            threshold = 0.5,
            )
            docs_question = retriever.get_relevant_documents(question)
            
            # Rerank
            if reranker is not None:
                with suppress_output():
                    docs_question = rerank_docs(reranker,docs_question,question)
            else:
                # Add a default reranking score
                for doc in docs_question:
                    doc.metadata["reranking_score"] = doc.metadata["similarity_score"]
                
            # If rerank by question we select the top documents for each question
            if rerank_by_question:
                docs_question = docs_question[:k_by_question[i]]
                
            # Add sources used in the metadata
            for doc in docs_question:
                doc.metadata["sources_used"] = sources
            
            # Add to the list of docs
            docs.extend(docs_question)
            
        # Sorting the list in descending order by rerank_score
        # Then select the top k
        docs = sorted(docs, key=lambda x: x.metadata["reranking_score"], reverse=True)
        docs = docs[:k_final]
        
        new_state = {"documents":docs}
        return new_state
    
    return retrieve_documents