AMG-RAG (Agentic Medical Graph-RAG) is a comprehensive framework that automates the construction and continuous updating of Medical Knowledge Graphs (MKGs), integrates reasoning, and retrieves current external evidence for medical Question Answering (QA). Our approach addresses the challenge of rapidly evolving medical knowledge by dynamically linking new findings and complex medical concepts.

• 🧠 Enhanced Knowledge Graph Construction: Advanced entity extraction with confidence scoring (1-10 scale)
• 🔄 Bidirectional Relationship Analysis: Comprehensive relationship mapping with confidence scoring
• 🎯 Context-Aware Entity Processing: LLM-generated descriptions with medical context integration
• 📚 Multi-source Evidence Retrieval: Integrates PubMed search, Wikipedia, and vector database retrieval
• 🔗 Chain-of-Thought Reasoning: Structured reasoning synthesis with evidence integration
• ⚡ Real-time Graph Updates: Dynamically incorporates latest medical literature and research
• 📊 Entity Summarization: Enhanced entity understanding with relevance-based confidence scoring

Our evaluations on standard medical QA benchmarks demonstrate superior performance:

AMG-RAG surpasses both comparable models and those 10 to 100 times larger, while enhancing interpretability for medical queries.

The enhanced AMG-RAG system consists of six key components:

• Structured output with relevance scoring (1-10 scale)
• Context-aware entity descriptions
• Confidence scoring based on relevance and external sources

• Bidirectional relationship analysis (A→B and B→A)
• Medical relationship types (treats, causes, symptom_of, risk_factor_for, etc.)
• Evidence-based confidence scoring

• PubMed API integration for latest research
• Wikipedia fallback for additional context
• Vector database semantic search

• LLM-generated enhanced summaries
• Relevance-based confidence updates
• Context integration for better understanding

• Structured reasoning synthesis with evidence integration
• Graph-based path exploration
• Confidence propagation through reasoning chains

• Multi-evidence integration for answer selection
• Confidence scoring and explanation generation

• Python 3.8+
• OpenAI API key (or Ollama for local inference)
• PubMed API key (optional, for higher rate limits)

# Clone the repository
git clone https://github.com/MrRezaeiUofT/AMG-RAG.git
cd AMG-RAG

# Install core dependencies
pip install langchain langchain-community langchain-openai
pip install transformers langgraph pandas numpy requests wikipedia networkx python-decouple

# Optional: Install additional packages
pip install langchain-huggingface langchain-chroma langchain-ollama

Create a .env file in the root directory:

OPENAI_API_KEY=your_openai_api_key_here
pubmed_api=your_pubmed_api_key_here

from AMG_with_KG import AMG_RAG_System

# Initialize the system
system = AMG_RAG_System(use_openai=True, openai_key="your-api-key")

# Sample medical question
question_data = {
    "question": "A 45-year-old man presents with severe chest pain...",
    "options": {
        "A": "Unstable angina",
        "B": "Acute inferior wall myocardial infarction",
        "C": "Acute anterior wall myocardial infarction",
        "D": "Aortic dissection",
        "E": "Pulmonary embolism"
    },
    "answer": "B"
}

# Get answer with reasoning
result = system.answer_question(question_data)

print(f"Answer: {result['answer']}")
print(f"Confidence: {result['confidence']:.2f}")
print(f"Explanation: {result['explanation']}")

# Access the knowledge graph
kg = system.kg

# Explore entities
for entity_name, entity in kg.entities.items():
    print(f"Entity: {entity_name}")
    print(f"Type: {entity.entity_type}")
    print(f"Confidence: {entity.confidence:.2f}")
    print(f"Description: {entity.description[:100]}...")

# Explore relationships
for relation in kg.relations:
    print(f"{relation.source} --[{relation.relation_type}]--> {relation.target}")
    print(f"Confidence: {relation.confidence:.2f}")

The system expects input data in JSONL format:

{
  "question": "Medical question text",
  "options": {
    "A": "Option A text",
    "B": "Option B text", 
    "C": "Option C text",
    "D": "Option D text"
  },
  "answer": "B",
  "answer_idx": 1,
  "meta_info": "Additional metadata"
}

# OpenAI (recommended)
system = AMG_RAG_System(use_openai=True, openai_key="your-api-key")

# Local Ollama
system = AMG_RAG_System(use_openai=False)

# Entity extraction settings
max_entities = 8              # Maximum entities per question
relevance_threshold = 5       # Minimum relevance score (1-10)

# Relationship analysis settings
confidence_threshold = 0.3    # Minimum confidence for relationships
max_relationship_depth = 2    # Maximum exploration depth

# Search parameters
pubmed_max_results = 3        # Max PubMed articles
wikipedia_sentences = 3       # Wikipedia summary length

AMG-RAG/
├── AMG-with-KG.py          # Main AMG-RAG system
├── Simple_AMG_RAG.py       # Simplified version
├── create_VDB.py           # Vector database utilities
├── dataset/                # Input datasets
│   ├── MEDQA/             # MEDQA dataset
│   ├── MedMCQA/           # MedMCQA dataset
│   └── PubMedQA/          # PubMedQA dataset
├── results/                # Output results
├── new_VDB/               # Vector database storage
├── requirements.txt        # Dependencies
├── .env                   # Environment variables
└── README.md             # This file

• 🎯 Relevance Scoring: 1-10 scale for entity importance
• 🔄 Bidirectional Relationships: Comprehensive A→B and B→A analysis
• 🌐 Context Integration: PubMed and Wikipedia context for better understanding
• 📝 Entity Summarization: LLM-generated enhanced descriptions

If you use AMG-RAG in your research, please cite our paper:

@misc{rezaei2025agenticmedicalknowledgegraphs,
      title={Agentic Medical Knowledge Graphs Enhance Medical Question Answering: Bridging the Gap Between LLMs and Evolving Medical Knowledge}, 
      author={Mohammad Reza Rezaei and Reza Saadati Fard and Jayson L. Parker and Rahul G. Krishnan and Milad Lankarany},
      year={2025},
      eprint={2502.13010},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2502.13010}, 
}

This project is licensed under the Apache-2.0 License - see the LICENSE file for details.

• Built with LangChain and LangGraph
• Uses Hugging Face Transformers for embeddings
• Integrates PubMed API for medical literature retrieval
• Enhanced with NetworkX for knowledge graph operations
• Benchmarked on MEDQA and MEDMCQA datasets

For questions, issues, or support:

1. Check the Issues page
2. Create a new issue with detailed information
3. Contact the maintainers

⚠️ Note: This is research software. Please validate results thoroughly before any clinical application.