knowledge_graph.py
1 """ 2 Graph-Assisted RAG: Relationship-Aware Retrieval 3 4 This example demonstrates how graph structures can enhance RAG by 5 capturing relationships between entities and enabling path-based retrieval. 6 7 RAG Concept: Traditional vector search finds similar documents but misses 8 relationships. Graph-assisted RAG understands connections between entities, 9 enabling multi-hop reasoning and relationship-aware answers. 10 """ 11 12 from praisonaiagents import Agent 13 14 # Sample knowledge base: Interconnected entities (simulating a knowledge graph) 15 ENTITY_KNOWLEDGE = [ 16 { 17 "id": "person_einstein", 18 "type": "person", 19 "content": """ 20 Albert Einstein (1879-1955) was a theoretical physicist who developed 21 the theory of relativity. He received the Nobel Prize in Physics in 1921 22 for his explanation of the photoelectric effect. Einstein worked at the 23 Institute for Advanced Study in Princeton from 1933 until his death. 24 Key relationships: Collaborated with Niels Bohr, influenced by Max Planck, 25 mentored by Heinrich Weber. 26 """ 27 }, 28 { 29 "id": "person_bohr", 30 "type": "person", 31 "content": """ 32 Niels Bohr (1885-1962) was a Danish physicist who made foundational 33 contributions to understanding atomic structure and quantum theory. 34 He received the Nobel Prize in Physics in 1922. Bohr founded the 35 Institute of Theoretical Physics in Copenhagen. 36 Key relationships: Debated with Einstein on quantum mechanics, 37 mentored Werner Heisenberg, influenced by Ernest Rutherford. 38 """ 39 }, 40 { 41 "id": "person_heisenberg", 42 "type": "person", 43 "content": """ 44 Werner Heisenberg (1901-1976) was a German physicist and pioneer of 45 quantum mechanics. He formulated the uncertainty principle in 1927. 46 He received the Nobel Prize in Physics in 1932. 47 Key relationships: Student of Niels Bohr, collaborated with Max Born, 48 worked with Wolfgang Pauli. 49 """ 50 }, 51 { 52 "id": "concept_relativity", 53 "type": "concept", 54 "content": """ 55 Theory of Relativity: A framework developed by Albert Einstein consisting 56 of special relativity (1905) and general relativity (1915). Special 57 relativity deals with objects moving at constant speeds, introducing 58 E=mc². General relativity describes gravity as the curvature of spacetime. 59 Related concepts: Spacetime, time dilation, gravitational waves. 60 Applications: GPS satellites, particle accelerators, cosmology. 61 """ 62 }, 63 { 64 "id": "concept_quantum", 65 "type": "concept", 66 "content": """ 67 Quantum Mechanics: A fundamental theory describing nature at atomic scales. 68 Key principles include wave-particle duality, superposition, and entanglement. 69 Developed through contributions from Planck, Bohr, Heisenberg, Schrödinger. 70 Related concepts: Uncertainty principle, quantum tunneling, wave function. 71 Applications: Semiconductors, lasers, MRI machines, quantum computing. 72 """ 73 }, 74 { 75 "id": "institution_princeton", 76 "type": "institution", 77 "content": """ 78 Institute for Advanced Study (Princeton): A private research center 79 founded in 1930. Notable members include Albert Einstein, John von Neumann, 80 Kurt Gödel, and Robert Oppenheimer. The institute focuses on theoretical 81 research in mathematics, natural sciences, and humanities. 82 Location: Princeton, New Jersey, USA. 83 """ 84 }, 85 { 86 "id": "institution_copenhagen", 87 "type": "institution", 88 "content": """ 89 Niels Bohr Institute (Copenhagen): Founded in 1921 by Niels Bohr as the 90 Institute of Theoretical Physics. It became a world center for quantum 91 mechanics research. Notable visitors included Heisenberg, Pauli, and Dirac. 92 The "Copenhagen interpretation" of quantum mechanics originated here. 93 Location: Copenhagen, Denmark. 94 """ 95 }, 96 { 97 "id": "event_solvay", 98 "type": "event", 99 "content": """ 100 Solvay Conferences: Series of physics conferences held in Brussels since 1911. 101 The 1927 Fifth Solvay Conference is famous for the Bohr-Einstein debates 102 on quantum mechanics. Attendees included Einstein, Bohr, Heisenberg, 103 Schrödinger, Dirac, and Marie Curie. These conferences shaped modern physics. 104 """ 105 } 106 ] 107 108 109 def graph_aware_rag(): 110 """Demonstrate relationship-aware retrieval.""" 111 112 # Build context from entity knowledge 113 context = "\n\n".join([f"[{e['type'].upper()}: {e['id']}]\n{e['content']}" for e in ENTITY_KNOWLEDGE]) 114 115 # Create agent with entity knowledge 116 agent = Agent( 117 name="Physics Historian", 118 instructions=f"""You are a physics historian who understands the relationships 119 between scientists, concepts, and institutions. When answering: 120 1. Trace connections between entities 121 2. Explain how ideas and people influenced each other 122 3. Reference specific relationships mentioned in the knowledge base 123 4. Build a narrative that shows the interconnected nature of physics history. 124 125 KNOWLEDGE GRAPH: 126 {context}""", 127 output="silent" 128 ) 129 130 # Relationship-focused queries 131 queries = [ 132 "How were Einstein and Bohr connected?", 133 "Trace the lineage of quantum mechanics through its key figures.", 134 "What institutions were central to 20th century physics?", 135 "How did the Solvay Conferences shape physics?" 136 ] 137 138 print("=" * 60) 139 print("GRAPH-ASSISTED RAG: Physics History Network") 140 print("=" * 60) 141 print("\nKnowledge graph entities:") 142 for entity in ENTITY_KNOWLEDGE: 143 print(f" - [{entity['type']}] {entity['id']}") 144 145 for query in queries: 146 print(f"\n📝 Query: {query}") 147 response = agent.chat(query) 148 print(f"💡 Answer: {response[:350]}..." if len(str(response)) > 350 else f"💡 Answer: {response}") 149 print("-" * 40) 150 151 152 def multi_hop_reasoning(): 153 """Demonstrate multi-hop reasoning across entities.""" 154 155 print("\n" + "=" * 60) 156 print("MULTI-HOP REASONING") 157 print("=" * 60) 158 159 # Build context 160 context = "\n\n".join([f"[{e['type'].upper()}: {e['id']}]\n{e['content']}" for e in ENTITY_KNOWLEDGE]) 161 162 agent = Agent( 163 name="Connection Finder", 164 instructions=f"""You excel at finding indirect connections between entities. 165 When asked about relationships, trace the path through intermediate entities. 166 Format your answer to show the chain of connections. 167 168 KNOWLEDGE GRAPH: 169 {context}""", 170 output="silent" 171 ) 172 173 # Multi-hop query requiring traversal 174 query = """ 175 What is the connection between the uncertainty principle and Princeton? 176 (Hint: trace through people and their relationships) 177 """ 178 179 print(f"\n📝 Multi-hop Query: {query.strip()}") 180 print("\n🔗 Expected path: Uncertainty Principle → Heisenberg → Bohr → Einstein → Princeton") 181 182 response = agent.chat(query) 183 print(f"\n💡 Answer:\n{response}") 184 185 186 def entity_type_filtering(): 187 """Show filtering by entity type.""" 188 189 print("\n" + "=" * 60) 190 print("ENTITY TYPE FILTERING") 191 print("=" * 60) 192 193 # Filter to only people 194 people_only = [e for e in ENTITY_KNOWLEDGE if e['type'] == 'person'] 195 people_context = "\n\n".join([f"[{e['id']}]\n{e['content']}" for e in people_only]) 196 197 people_agent = Agent( 198 name="Biographer", 199 instructions=f"""You are a biographer focused on scientists' lives and relationships. 200 201 SCIENTISTS: 202 {people_context}""", 203 output="silent" 204 ) 205 206 # Filter to only concepts 207 concepts_only = [e for e in ENTITY_KNOWLEDGE if e['type'] == 'concept'] 208 concepts_context = "\n\n".join([f"[{e['id']}]\n{e['content']}" for e in concepts_only]) 209 210 concept_agent = Agent( 211 name="Concept Explainer", 212 instructions=f"""You explain scientific concepts and their relationships. 213 214 CONCEPTS: 215 {concepts_context}""", 216 output="silent" 217 ) 218 219 query = "Tell me about the major developments in early 20th century physics." 220 221 print(f"\n📝 Query: {query}") 222 223 print("\n👤 People-focused view:") 224 response1 = people_agent.chat(query) 225 print(f" {response1[:250]}...") 226 227 print("\n💡 Concept-focused view:") 228 response2 = concept_agent.chat(query) 229 print(f" {response2[:250]}...") 230 231 232 def graph_rag_patterns(): 233 """Explain graph RAG patterns.""" 234 235 print("\n" + "=" * 60) 236 print("GRAPH RAG PATTERNS") 237 print("=" * 60) 238 239 print(""" 240 📊 Graph-Assisted RAG Patterns: 241 242 1. **Entity-Centric Retrieval** 243 - Store entities with their relationships 244 - Query retrieves entity + connected entities 245 - Good for: "Tell me about X and its connections" 246 247 2. **Path-Based Retrieval** 248 - Find paths between two entities 249 - Retrieve all nodes along the path 250 - Good for: "How is X related to Y?" 251 252 3. **Neighborhood Expansion** 253 - Start with query-matched entities 254 - Expand to include 1-hop or 2-hop neighbors 255 - Good for: Comprehensive context 256 257 4. **Subgraph Extraction** 258 - Extract relevant subgraph for query 259 - Include nodes and edges as context 260 - Good for: Complex multi-entity questions 261 262 Implementation in PraisonAI: 263 ```python 264 # Structure knowledge as entities with relationships 265 entities = [ 266 { 267 "id": "entity_1", 268 "type": "person", 269 "content": "...", 270 "relationships": ["entity_2", "entity_3"] 271 }, 272 ... 273 ] 274 275 # Agent retrieves related entities automatically 276 agent = Agent( 277 knowledge=entities, 278 instructions="Trace relationships between entities..." 279 ) 280 ``` 281 """) 282 283 284 def main(): 285 """Run all graph-assisted RAG examples.""" 286 print("\n🚀 PraisonAI Graph-Assisted RAG Examples\n") 287 288 # Example 1: Graph-aware retrieval 289 graph_aware_rag() 290 291 # Example 2: Multi-hop reasoning 292 multi_hop_reasoning() 293 294 # Example 3: Entity type filtering 295 entity_type_filtering() 296 297 # Example 4: Graph RAG patterns 298 graph_rag_patterns() 299 300 print("\n✅ Graph-assisted RAG examples completed!") 301 302 303 if __name__ == "__main__": 304 main()