reflection.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This module contains the reflection logic for the RAG server.
Reflection is a technique used to improve the quality of the generated response by checking the relevance of the retrieved context and the groundedness of the generated response.
1. ReflectionCounter: A class that tracks the number of reflection iterations across query rewrites and response regeneration.
2. check_context_relevance: Check relevance of retrieved context and optionally rewrite query for better results.
3. check_response_groundedness: Check groundedness of generated response against retrieved context.
4. _retry_score_generation: Helper method to retry score generation with error handling.
"""

import logging
import os
from typing import List, Tuple, Dict, Any
from concurrent.futures import ThreadPoolExecutor

from langchain_core.output_parsers.string import StrOutputParser
from langchain_core.prompts.chat import ChatPromptTemplate
from langchain_core.runnables import RunnableAssign

from nvidia_rag.utils.llm import get_llm, get_prompts
from nvidia_rag.utils.common import get_env_variable
from nvidia_rag.utils.vectorstore import retreive_docs_from_retriever

logger = logging.getLogger(__name__)
prompts = get_prompts()

def _retry_score_generation(chain, inputs: Dict[str, Any], max_retries: int = 3, config: Dict[str, Any] = {}) -> int:
    """Helper method to retry score generation with error handling.

    Args:
        chain: The chain to execute
        inputs: Input dictionary for the chain
        max_retries: Maximum number of retry attempts

    Returns:
        int: Generated score (0, 1, or 2), or 0 if all retries fail
    """
    for retry in range(max_retries):
        try:
            response = chain.invoke(inputs, config=config)
            # Extract numeric score from response
            for score in [2, 1, 0]:
                if str(score) in response:
                    return score
        except Exception as e:
            logger.warning(f"Retry {retry + 1}/{max_retries} failed: {str(e)}")
            if retry == max_retries - 1:
                logger.error(f"All retries failed for score generation")
                return 0
            continue
    return 0

class ReflectionCounter:
    """Tracks the number of reflection iterations across query rewrites and response regeneration."""
    def __init__(self, max_loops: int):
        self.max_loops = max_loops
        self.current_count = 0

    def increment(self) -> bool:
        """Increment counter and return whether we can continue."""
        if self.current_count >= self.max_loops:
            return False
        self.current_count += 1
        return True

    @property
    def remaining(self) -> int:
        return max(0, self.max_loops - self.current_count)

def check_context_relevance(retriever_query: str,
                          retrievers: List[Any],
                          ranker,
                          reflection_counter: ReflectionCounter,
                          enable_reranker: bool = True,
                          filter_expr: str = ''
                          ) -> Tuple[List[str], bool]:
    """Check relevance of retrieved context and optionally rewrite query for better results.

    Args:
        retriever_query (str): Current query to use for retrieval
        retrievers: List of Document retriever instances
        ranker: Optional document ranker instance
        reflection_counter: ReflectionCounter instance to track loop count
        enable_reranker: Whether to use the reranker if available
        filter_expr: Filter expression to filter the retrieved documents from Milvus collection

    Returns:
        Tuple[List[str], bool]: Retrieved documents and whether they meet relevance threshold
    """
    relevance_threshold = int(os.environ.get("CONTEXT_RELEVANCE_THRESHOLD", 1))
    reflection_llm_name = get_env_variable(variable_name="REFLECTION_LLM", default_value="mistralai/mixtral-8x22b-instruct-v0.1").strip('"').strip("'")
    reflection_llm_endpoint = os.environ.get("REFLECTION_LLM_SERVERURL", "").strip('"').strip("'")

    llm_params = {
        "model": reflection_llm_name,
        "temperature": 0.2,
        "top_p": 0.9,
        "max_tokens": 512
    }

    if reflection_llm_endpoint:
        llm_params["llm_endpoint"] = reflection_llm_endpoint

    reflection_llm = get_llm(**llm_params)

    relevance_template = ChatPromptTemplate.from_messages([
        ("system", prompts["reflection_relevance_check_prompt"]["system"]),
        ("human", "{query}\n\n{context}")
    ])

    query_rewrite_template = ChatPromptTemplate.from_messages([
        ("system", prompts["reflection_query_rewriter_prompt"]["system"]),
        ("human", "{query}")
    ])

    current_query = retriever_query

    while reflection_counter.remaining > 0:
        # Get documents using current query
        if ranker and enable_reranker:
            context_reranker = RunnableAssign({
                "context":
                    lambda input: ranker.compress_documents(query=input['question'], documents=input['context'])
            })

            # Perform parallel retrieval from all vector stores
            docs = []
            with ThreadPoolExecutor() as executor:
                futures = [executor.submit(retreive_docs_from_retriever, retriever=retriever, retriever_query=current_query, expr=filter_expr) for retriever in retrievers]
                for future in futures:
                    docs.extend(future.result())

            docs = context_reranker.invoke({"context": docs, "question": current_query}, config={'run_name':'context_reranker'})
            original_docs = docs.get("context", [])
        else:
            # Perform sequential retrieval from the first vector store
            original_docs = retreive_docs_from_retriever(retriever=retrievers[0], retriever_query=current_query, expr=filter_expr)

        docs = [d.page_content for d in original_docs]

        context_text = "\n".join(docs)
        relevance_chain = relevance_template | reflection_llm | StrOutputParser()
        relevance_score = _retry_score_generation(
            relevance_chain,
            {"query": current_query, "context": context_text},
            config={'run_name':'relevance-checker'}
        )

        logger.info(f"Context relevance score: {relevance_score} (threshold: {relevance_threshold})")
        reflection_counter.increment()

        if relevance_score >= relevance_threshold:
            return original_docs, True

        if reflection_counter.remaining > 0:
            rewrite_chain = query_rewrite_template | reflection_llm | StrOutputParser()
            current_query = rewrite_chain.invoke({"query": current_query}, config={'run_name':'query-rewriter'})
            logger.info(f"Rewritten query (iteration {reflection_counter.current_count}): {current_query}")

    return original_docs, False

def check_response_groundedness(response: str,
                              context: List[str],
                              reflection_counter: ReflectionCounter,
                              ) -> Tuple[str, bool]:
    """Check groundedness of generated response against retrieved context.

    Args:
        response (str): Generated response to check
        context (List[str]): List of context documents
        reflection_counter: ReflectionCounter instance to track loop count

    Returns:
        Tuple[str, bool]: Final response and whether it meets groundedness threshold
    """
    groundedness_threshold = int(os.environ.get("RESPONSE_GROUNDEDNESS_THRESHOLD", 1))
    reflection_llm_name = get_env_variable(variable_name="REFLECTION_LLM", default_value="mistralai/mixtral-8x22b-instruct-v0.1").strip('"').strip("'")
    reflection_llm_endpoint = os.environ.get("REFLECTION_LLM_SERVERURL", "").strip('"').strip("'")

    llm_params = {
        "model": reflection_llm_name,
        "temperature": 0.2,
        "top_p": 0.9,
        "max_tokens": 1024
    }

    if reflection_llm_endpoint:
        llm_params["llm_endpoint"] = reflection_llm_endpoint

    reflection_llm = get_llm(**llm_params)

    groundedness_template = ChatPromptTemplate.from_messages([
        ("system", prompts["reflection_groundedness_check_prompt"]["system"]),
        ("human", "{context}\n\n{response}")
    ])

    context_text = "\n".join(context)
    current_response = response

    while reflection_counter.remaining > 0:
        groundedness_chain = groundedness_template | reflection_llm | StrOutputParser()
        groundedness_score = _retry_score_generation(
            groundedness_chain,
            {"context": context_text, "response": current_response}
        )

        logger.info(f"Response groundedness score: {groundedness_score} (threshold: {groundedness_threshold})")
        reflection_counter.increment()

        if groundedness_score >= groundedness_threshold:
            return current_response, True

        if reflection_counter.remaining > 0:
            regen_prompt = ChatPromptTemplate.from_messages([
                ("system", prompts["reflection_response_regeneration_prompt"]["system"]),
                ("human", f"Context: {context_text}\n\nPrevious response: {current_response}\n\n"
                         "Generate a new, more grounded response:")
            ])

            regen_chain = regen_prompt | reflection_llm | StrOutputParser()
            current_response = regen_chain.invoke({}, config={'run_name':'response-regenerator'})
            logger.info(f"Regenerated response (iteration {reflection_counter.current_count})")

    return current_response, False