Building a Hallucination Checker

import os
import re
import json
from typing import Dict, Any, List, Annotated
from pydantic import BaseModel
from langchain_core.tools import StructuredTool
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain_exa import ExaSearchRetriever
from langchain_core.runnables import RunnableLambda
from langchain_core.prompts import PromptTemplate
from langchain_anthropic import ChatAnthropic

# Check for API keys
assert os.getenv("EXA_API_KEY"), "Please set the EXA_API_KEY environment variable"
assert os.getenv("ANTHROPIC_API_KEY"), "Please set the ANTHROPIC_API_KEY environment variable"

# Set up the LLM (ChatAnthropic)
llm = ChatAnthropic(model="claude-3-5-sonnet-20240620", temperature=0)

def extract_claims_regex(text: str) -> List[str]:
    """Fallback function to extract claims using regex."""
    pattern = r'([A-Z][^.!?]*?[.!?])'
    matches = re.findall(pattern, text)
    return [match.strip()+'.' for match in matches]

def extract_claims(text: str) -> List[str]:
    """Extract factual claims from the text using an LLM."""
    system_message = SystemMessage(content="""
    You are an expert at extracting claims from text.
    Your task is to identify and list all claims present, true or false,
    in the given text. Each claim should be a single, verifiable statement.
    Consider various forms of claims, including assertions, statistics, and
    quotes. Do not skip any claims, even if they seem obvious. Do not include in the list 'The text contains a claim that needs to be checked for hallucinations' - this is not a claim.
    Present the claims as a JSON array of strings, and do not include any additional text.
    """)

    human_message = HumanMessage(content=f"Extract factual claims from this text: {text}")
    response = llm.invoke([system_message, human_message])

    try:
        claims = json.loads(response.content)
        if not isinstance(claims, list):
            raise ValueError("Response is not a list")
    except (json.JSONDecodeError, ValueError):
        # Fallback to regex extraction if LLM response is not valid JSON
        claims = extract_claims_regex(text)
    
    return claims

def exa_search(query: str) -> List[str]:
    """Function to retrieve usable documents for AI assistant."""
    search = ExaSearchRetriever(k=5, text=True)

    print("Query: ", query)

    document_prompt = PromptTemplate.from_template(
        """
        <source>
            <url>{url}</url>
            <text>{text}</text>
        </source>
        """
    )

    parse_info = RunnableLambda(
        lambda document: {
            "url": document.metadata["url"],
            "text": document.page_content or "No text available",
        }
    )

    document_chain = (parse_info | document_prompt)
    search_chain = search | document_chain.map()
    documents = search_chain.invoke(query+".\n Here is a web page to help verify this claim:")

    print("Documents: ", documents)
    
    return [str(doc) for doc in documents]

def verify_claim(claim: str, sources: List[str]) -> Dict[str, Any]:
    """Verify a single claim using combined Exa search sources."""
    if not sources:
        # If no sources are returned, default to insufficient information
        return {
            "claim": claim,
            "assessment": "Insufficient information",
            "confidence_score": 0.5,
            "supporting_sources": [],
            "refuting_sources": []
        }
    
    # Combine the sources into one text
    combined_sources = "\n\n".join(sources)
    
    system_message = SystemMessage(content="""
    You are an expert fact-checker.
    Given a claim and a set of sources, determine whether the claim is supported, refuted, or if there is insufficient information in the sources to make a determination.
    For your analysis, consider all the sources collectively.
    Provide your answer as a JSON object with the following structure:
    {
        "claim": "...",
        "assessment": "supported" or "refuted" or "Insufficient information",
        "confidence_score": a number between 0 and 1 (1 means fully confident the claim is true, 0 means fully confident the claim is false),
        "supporting_sources": [list of sources that support the claim],
        "refuting_sources": [list of sources that refute the claim]
    }
    Do not include any additional text.
    """)
    
    human_message = HumanMessage(content=f"""
    Claim: "{claim}"
    
    Sources:
    {combined_sources}
    
    Based on the above sources, assess the claim.
    """)
    
    response = llm.invoke([system_message, human_message])
    
    try:
        result = json.loads(response.content)
        if not isinstance(result, dict):
            raise ValueError("Response is not a JSON object")
    except (json.JSONDecodeError, ValueError):
        # If parsing fails, default to insufficient information
        result = {
            "claim": claim,
            "assessment": "Insufficient information",
            "confidence_score": 0.5,
            "supporting_sources": [],
            "refuting_sources": []
        }
    
    return result

def hallucination_check(text: str) -> Dict[str, Any]:
    """Check a given text for hallucinations using Exa search."""
    claims = extract_claims(text)
    claim_verifications = []

    for claim in claims:
        sources = exa_search(claim)
        verification_result = verify_claim(claim, sources)
        claim_verifications.append(verification_result)

    return {
        "claims": claim_verifications
    }

def hallucination_check_tool(text: str) -> Dict[str, Any]:
    """Assess the given text for hallucinations using Exa search."""
    return hallucination_check(text)

structured_tool = StructuredTool.from_function(
    func=hallucination_check_tool,
    name="hallucination_check",
    description="Assess the given text for hallucinations using Exa search."
)

class State(BaseModel):
    messages: Annotated[List, add_messages]
    analysis_result: Dict[str, Any] = {}

def call_model(state: State):
    # Simulate the assistant calling the tool
    return {"messages": state.messages + [AIMessage(content="Use hallucination_check tool", additional_kwargs={"tool_calls": [{"type": "function", "function": {"name": "hallucination_check"}}]})]}

def run_tool(state: State):
    text_to_check = next((m.content for m in reversed(state.messages) if isinstance(m, HumanMessage)), "")
    tool_output = structured_tool.invoke(text_to_check)
    return {"messages": state.messages + [AIMessage(content=str(tool_output))], "analysis_result": tool_output}

def use_analysis(state: State) -> str:
    return "tools"

workflow = StateGraph(State)
workflow.add_node("agent", call_model)
workflow.add_node("tools", run_tool)
workflow.add_node("process_result", lambda x: x)
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", use_analysis, {
    "tools": "tools"
})
workflow.add_edge("tools", "process_result")
workflow.add_edge("process_result", END)

graph = workflow.compile()

initial_state = State(messages=[
    SystemMessage(content="You are a helpful assistant."),
    HumanMessage(content="Check this text for hallucinations: The Eiffel Tower, an iconic iron lattice structure located in Paris, was originally constructed as a giant sundial in 1822.")
])

final_state = graph.invoke(initial_state)