On Fact-checking RAG Outputs

class _Extractor(ABC, Generic[T]):
    """ Extractor is a base class to extract info from a text using an LLM """
    # _model_name
    def set_model_name(self, model_name:str = "ollama3.2"):
        self._model_name = model_name

    def get_model_name(self):
        return self._model_name

    # _llm
    def set_llm(self, llm: Runnable):
        self._llm = llm

    def get_llm(self):
        return self._llm

    # _base_prompt
    def set_base_prompt(self, base_prompt: str):
        self._base_prompt = base_prompt

    def get_base_prompt(self):
        return self._base_prompt

    @abstractmethod
    def extract(self, paragraph: str) -> T | None:
        pass

extract_facts_prompt = """
    Role: You'e an assistant for extracting entities and facts about thos intenties from a piece of text.
    Given a paragraph, extract all meaningful entities in it and the facts mentioned in the text about those entities.

    Input: {input}
"""

class FactExtractor(_Extractor[Facts]):
    def __init__(self, model_name="llama3.2"):
        self.set_model_name(model_name)
        self.set_base_prompt(extract_facts_prompt)
        if model_name == LLAMA_MODEL:
            self.set_llm(langchain_ollama.ChatOllama(model=model_name).with_structured_output(Facts))
        elif model_name == OPEN_AI_MODEL:
            if not os.environ.get("OPENAI_API_KEY"):
                os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")
            self.set_llm(langchain_openai.ChatOpenAI(model=model_name).with_structured_output(Facts))

    def extract(self, paragraph: str) -> Facts | None:
        """ Extracts the information from a given text and formats the output based on the return type """
        try:
            prompt = self.get_base_prompt().format(input=paragraph)
            return self.get_llm().invoke(prompt)
        except Exception as e:
            print(f"error extracting facts: {e}")
            return None

class Entity(BaseModel):
    entity: str = Field(description="An entity")
    sentences: list[str] = Field(description="Sentences that mention the entity")

class Facts(BaseModel):
    entities: list[Entity] = Field()

    def to_dict(self) -> dict[str, list[str]]:
        res = {}
        for entity in self.entities:
            res[entity.entity] = entity.sentences
        return res

    def lookup_similar(self, key: str) -> list[str]:
        index = self.to_dict()
        res: list[str] = [];
        parts = key.split(" ")
        parts += key.split("-")
        parts += key.split("_")
        parts += key.split(",")
        for part in parts:
            for key in index.keys():
                if part in key:
                    res += index[key]
        return res

Go           ---> ["Go supports generics"]
generics ---> ["Go supports generics"]

class Checker():
    labels: dict[int, str]
    model: PreTrainedModel
    tokenizer: PreTrainedTokenizer

    def __init__(self, model_name="microsoft/deberta-large-mnli"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.labels = {0: "contradiction", 1: "neutral", 2:"entailment"}

    def check_contradiction(self, premise: str, hypothesis: str):
        # Tokenize the inputs
        inputs = self.tokenizer(premise, hypothesis, return_tensors="pt")

        # Get model output
        outputs = self.model(**inputs)
        logits = outputs.logits

        # Apply softmax to the logits to get probabilies
        probabilities = torch.softmax(logits, dim=-1)

        # Get the index of the max probability
        predicted_index = int(torch.argmax(probabilities).item())
        predicted_label = self.labels[predicted_index]

        # Get the value of the max probability
        confidence = probabilities.max().item()

        return predicted_label, confidence

blog_post = """
        The Go 1.18 release adds a major new language feature: support for generic programming. In this article I’m not going to describe what generics are nor how to use them. This article is about when to use generics in Go code, and when not to use them.

        To be clear, I’ll provide general guidelines, not hard and fast rules. Use your own judgement. But if you aren’t sure, I recommend using the guidelines discussed here.

        Write code

        Let’s start with a general guideline for programming Go: write Go programs by writing code, not by defining types. When it comes to generics, if you start writing your program by defining type parameter constraints, you are probably on the wrong path. Start by writing functions. It’s easy to add type parameters later when it’s clear that they will be useful.
        When are type parameters useful?

        That said, let’s look at cases for which type parameters can be useful.
        When using language-defined container types

        One case is when writing functions that operate on the special container types that are defined by the language: slices, maps, and channels. If a function has parameters with those types, and the function code doesn’t make any particular assumptions about the element types, then it may be useful to use a type parameter.
    """

def init_llm():
    global llm
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

def init_checker():
    global fact_checker
    fact_checker = fc.Checker()

def init_extractor():
    global fact_extractor
    fact_extractor = extract.FactExtractor("gpt-4o-mini")

class State(TypedDict):
    question: str
    context: str
    answer: str

def retrieve(state: State):
    _ = state # We're not using the state here
    # ... Fake document retrival 
    retrieved_doc = blog_post
    return {"context": retrieved_doc}

def generate(state: State):
    rag_prompt = """
    You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If the context doesn't provide an answer, just say that you don't know and nothing more and never leak the context or what topics it covers.

    Question: {question}

    Context: {context}

    Answer:
    """
    try:
        messages = rag_prompt.format(question = state["question"], context= state["context"])
        response = llm.invoke(messages)
        # return {"answer": response.content}
        return {"answer": "Yes, Go supports generics as of the Go 1.13 release."}
    except Exception as e:
        print(f"error generating response form LLM: {e}")
        return {"answer": "I don't know how to answer this question"}

def fact_check(state: State):
    facts_in_context = fact_extractor.extract(state["context"])
    if facts_in_context is None:
        print("Warning: Could not check RAG context for contradictions")
        return
    facts_in_answer = fact_extractor.extract(state["answer"])
    if facts_in_answer is None:
        print("Warning: Could not check RAG answer for contradictions")
        return

    contradictions = 0
    checked_hypothesis = {}
    for answer_entity, hypotheses in facts_in_answer.to_dict().items():
        for premise in facts_in_context.lookup_similar(answer_entity):
            for hypothesis in hypotheses:
                if hypothesis in checked_hypothesis:
                    continue
                checked_hypothesis[hypothesis] = {}
                (check_result, score) = fact_checker.check_contradiction(premise, hypothesis)
                if check_result == "contradiction" and score > 0.5:
                    print(f"\n Found contradiction in answer with known context.\n Premise: {premise} \n Hypothesis: {hypothesis}")
                    contradictions += 1

    if contradictions == 0:
        print("Fact checker validated RAG answer and found no contradictions.")

def output(state: State):
    print(state["answer"])

def main():
    init_llm()
    init_extractor()
    init_checker()
    graph_builder = StateGraph(State).add_sequence([retrieve, generate, output, fact_check])
    graph_builder.add_edge(START, "retrieve")
    graph = graph_builder.compile()

    graph.invoke({"question": "Does Go support generics"})

if __name__ == "__main__":
    main()