CommunityKG-RAG: Leveraging Community Structures in Knowledge Graphs for Advanced Retrieval-Augmented Generation in Fact-Checking

pip install networkx python-louvain openai

import os
import networkx as nx
import community
from openai import OpenAI
import dspy

# Set up OpenAI client
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Minimal Knowledge Graph implementation
class KnowledgeGraph:
    """
    Represents a simplified Knowledge Graph (KG) structure.
    In the paper, KGs are used to capture complex relationships between entities,
    which is crucial for providing rich context in fact-checking tasks.
    """
    def __init__(self):
        self.graph = nx.Graph()

    def add_entity(self, entity):
        """Add an entity (node) to the graph."""
        self.graph.add_node(entity)

    def add_relationship(self, entity1, entity2, relationship):
        """
        Add a relationship (edge) between two entities.
        This method captures the triple structure (subject, predicate, object) mentioned in the paper.
        """
        self.graph.add_edge(entity1, entity2, relationship=relationship)

    def get_community_structure(self):
        """
        Detect communities in the graph using the Louvain algorithm.
        This aligns with the paper's emphasis on leveraging community structures within KGs.
        """
        return community.best_partition(self.graph)

    def get_community_sentences(self, community_id, community_structure):
        """
        Retrieve sentences (nodes) belonging to a specific community.
        This method supports the community-based retrieval approach proposed in the paper.
        """
        return [node for node, com in community_structure.items() if com == community_id]

# CommunityKG-RAG implementation
class CommunityKGRAG:
    """
    Implements the core ideas of the CommunityKG-RAG approach.
    This class integrates community structures in KGs with RAG systems for enhanced fact-checking.
    """
    def __init__(self, kg):
        self.kg = kg
        # Precompute community structure, as suggested in the paper for efficiency
        self.community_structure = kg.get_community_structure()

    def retrieve_context(self, claim, top_k_communities=2, top_k_sentences=5):
        """
        Retrieve relevant context based on community structures.
        This method embodies the paper's novel approach of using community structures for context retrieval.
        
        Note: The current implementation uses a simplified word overlap method.
        The paper suggests using more sophisticated embedding models for similarity calculation.
        """
        claim_words = set(claim.lower().split())
        community_scores = {}
        
        # Score communities based on relevance to the claim
        for com_id in set(self.community_structure.values()):
            com_sentences = self.kg.get_community_sentences(com_id, self.community_structure)
            com_words = set(" ".join(com_sentences).lower().split())
            score = len(claim_words.intersection(com_words))
            community_scores[com_id] = score

        # Select top-k most relevant communities
        top_communities = sorted(community_scores, key=community_scores.get, reverse=True)[:top_k_communities]
        
        # Retrieve sentences from top communities
        context_sentences = []
        for com_id in top_communities:
            sentences = self.kg.get_community_sentences(com_id, self.community_structure)
            context_sentences.extend(sentences[:top_k_sentences])

        return " ".join(context_sentences)

    def verify_claim(self, claim):
        """
        Verify a claim using the CommunityKG-RAG approach.
        This method demonstrates the integration of KG-based retrieval with a language model for fact-checking.
        """
        # Retrieve context using community-based approach
        context = self.retrieve_context(claim)
        
        # Construct prompt for the language model
        prompt = f"""Given the evidence provided below:
{context}
Please evaluate the following claim:
{claim}
Based on the evidence, should the claim be rated as 'True', 'False', or 'NEI' (Not Enough Information)?"""

        # Use GPT-4 for claim verification, as suggested in the paper for its advanced reasoning capabilities
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}]
        )

        return response.choices[0].message.content

# Create and populate the Knowledge Graph
kg = KnowledgeGraph()

# Add example entities and relationships to demonstrate the KG structure
kg.add_entity("COVID-19")
kg.add_entity("Vaccine")
kg.add_entity("Pfizer")
kg.add_entity("Moderna")
kg.add_entity("mRNA")
kg.add_relationship("COVID-19", "Vaccine", "prevented by")
kg.add_relationship("Pfizer", "Vaccine", "produces")
kg.add_relationship("Moderna", "Vaccine", "produces")
kg.add_relationship("Vaccine", "mRNA", "uses")

# Add example sentences to the graph
# In a full implementation, these would be extracted from fact-checking articles as mentioned in the paper
kg.add_entity("COVID-19 vaccines have been shown to be effective in preventing severe illness.")
kg.add_entity("Pfizer and Moderna vaccines use mRNA technology.")
kg.add_entity("Vaccines undergo rigorous testing before approval.")
kg.add_entity("mRNA vaccines do not alter human DNA.")
kg.add_entity("Vaccine side effects are generally mild and short-lived.")

# Create the CommunityKG-RAG system
ckgrag = CommunityKGRAG(kg)

# Example claims to verify
claims = [
    "COVID-19 vaccines are effective in preventing severe illness.",
    "mRNA vaccines alter human DNA.",
    "Pfizer produces a COVID-19 vaccine."
]

# Verify the claims using the CommunityKG-RAG approach
for claim in claims:
    result = ckgrag.verify_claim(claim)
    print(f"Claim: {claim}")
    print(f"Verification result: {result}\n")

"""
Key aspects of the paper implemented in this code:

1. Knowledge Graph Structure: The KnowledgeGraph class represents the paper's emphasis on using KGs to capture complex entity relationships.

2. Community Detection: The get_community_structure method implements the paper's suggestion of using community detection (Louvain algorithm) to identify clusters of related information.

3. Community-Based Retrieval: The retrieve_context method in CommunityKGRAG class demonstrates the paper's novel approach of using community structures for context retrieval.

4. Integration with LLM: The verify_claim method shows how the retrieved context is used with a language model (GPT-4) for fact-checking, as proposed in the paper.

5. Zero-Shot Framework: The implementation doesn't require additional training, aligning with the paper's emphasis on a zero-shot approach.

Areas for further development (as mentioned in the paper):
- Implement more sophisticated embedding models for similarity calculation.
- Expand the knowledge graph with comprehensive data from fact-checking articles.
- Enhance the community detection and selection process.
- Improve the context retrieval mechanism to better utilize multi-hop relationships.

This prototype serves as a proof-of-concept for the CommunityKG-RAG approach, demonstrating its potential in enhancing fact-checking through the integration of structured knowledge and language models.
"""

python community-rag.py

Claim: COVID-19 vaccines are effective in preventing severe illness.
Verification result: Based on the evidence provided, the claim that "COVID-19 vaccines are effective in preventing severe illness" should be rated as **True**. The evidence explicitly states that COVID-19 vaccines have been shown to be effective in preventing severe illness, which directly supports the claim.

Claim: mRNA vaccines alter human DNA.
Verification result: Based on the evidence provided, the claim "mRNA vaccines alter human DNA" should be rated as 'False.' The evidence clearly states that mRNA vaccines do not alter human DNA, and since Pfizer and Moderna vaccines use mRNA technology, it follows that these specific vaccines also do not alter human DNA.

Claim: Pfizer produces a COVID-19 vaccine.
Verification result: Based on the evidence provided, the claim "Pfizer produces a COVID-19 vaccine" should be rated as 'True'. The evidence states that both the Pfizer and Moderna vaccines use mRNA technology, which confirms that Pfizer is indeed a producer of a COVID-19 vaccine.