File size: 37,588 Bytes
9e5bc69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
"""

Knowledge Retrieval Module - Phase C (Steps 6-8)



Performs community search and data extraction using graph database structures.

Handles community retrieval, data extraction, and initial answer generation.

"""

import logging
import numpy as np
import json
from typing import Dict, List, Tuple, Any
from dataclasses import dataclass
from datetime import datetime

from .setup import GraphRAGSetup
from .query_preprocessing import DriftRoutingResult


@dataclass
class CommunityResult:
    """Enhanced community result with comprehensive properties."""
    community_id: str
    similarity_score: float
    summary: str
    key_entities: List[str]
    member_ids: List[str]  # Direct member access
    modularity_score: float  # Community quality
    level: int
    internal_edges: int
    member_count: int
    centrality_stats: Dict[str, float]  # Aggregated centrality measures
    confidence_score: float
    search_index: str  # Optimized search key
    termination_criteria: Dict[str, Any]


@dataclass
class EntityResult:
    """Entity result with attributes from graph database."""
    entity_id: str
    name: str
    content: str
    confidence: float
    degree_centrality: float
    betweenness_centrality: float
    closeness_centrality: float
    community_id: str
    node_type: str


@dataclass
class RelationshipResult:
    """Relationship result with graph database attributes."""
    start_node: str
    end_node: str
    relationship_type: str
    confidence: float


class CommunitySearchEngine:
    """Knowledge retrieval engine for community search and entity extraction."""
    
    def __init__(self, setup: GraphRAGSetup):
        self.setup = setup
        self.neo4j_conn = setup.neo4j_conn
        self.config = setup.config
        self.logger = logging.getLogger(self.__class__.__name__)
        
        # Initialize search optimization
        self.community_search_index = {}
        self.centrality_cache = {}
        
    async def execute_primer_phase(self,

                                 query_embedding: List[float],

                                 routing_result: DriftRoutingResult) -> Dict[str, Any]:
        """Execute community search and knowledge retrieval."""
        start_time = datetime.now()
        
        try:
            # Community retrieval
            self.logger.info("Starting community retrieval")
            communities = await self._retrieve_communities_enhanced(
                query_embedding, routing_result
            )
            
            # Data extraction
            self.logger.info("Starting data extraction")
            extracted_data = await self._extract_community_data_enhanced(communities)
            
            # Answer generation
            self.logger.info("Starting answer generation")
            initial_answer = await self._generate_initial_answer_enhanced(
                extracted_data, routing_result
            )
            
            execution_time = (datetime.now() - start_time).total_seconds()
            
            return {
                'communities': communities,
                'extracted_data': extracted_data,
                'initial_answer': initial_answer,
                'execution_time': execution_time,
                'metadata': {
                    'communities_retrieved': len(communities),
                    'entities_extracted': len(extracted_data.get('entities', [])),
                    'relationships_extracted': len(extracted_data.get('relationships', [])),
                    'phase': 'primer',
                    'step_range': '6-8'
                }
            }
            
        except Exception as e:
            self.logger.error(f"Primer phase execution failed: {e}")
            raise
    
    async def _retrieve_communities_enhanced(self,

                                           query_embedding: List[float],

                                           routing_result: DriftRoutingResult) -> List[CommunityResult]:
        """

        Step 6: Enhanced community retrieval using comprehensive properties.

        

        Retrieves relevant communities based on query embedding similarity.

        """
        try:
            # Retrieve HyDE embeddings
            hyde_embeddings = await self._retrieve_hyde_embeddings_enhanced()
            
            if not hyde_embeddings:
                self.logger.warning("No HyDE embeddings found")
                return []
            
            # Compute similarities
            similarities = self._compute_hyde_similarities_enhanced(
                query_embedding, hyde_embeddings
            )
            
            # Rank communities
            ranked_communities = self._rank_communities_enhanced(
                similarities, routing_result
            )
            
            # Apply criteria
            filtered_communities = self._apply_termination_criteria(
                ranked_communities, routing_result
            )
            
            # Fetch community details
            community_results = await self._fetch_community_details_enhanced(
                filtered_communities
            )
            
            self.logger.info(f"Retrieved {len(community_results)} enhanced communities")
            return community_results
            
        except Exception as e:
            self.logger.error(f"Enhanced community retrieval failed: {e}")
            return []
    
    async def _load_community_search_index(self):
        """Load optimized community search index from Neo4j."""
        try:
            query = """

            MATCH (meta:DriftMetadata)

            WHERE meta.community_search_index IS NOT NULL

            RETURN meta.community_search_index as search_index,

                   meta.total_communities as total_communities

            """
            
            results = self.neo4j_conn.execute_query(query)
            
            for record in results:
                # The search index is a nested JSON structure with community IDs as keys
                search_index_data = record['search_index']
                if isinstance(search_index_data, dict):
                    # Each community in the search index 
                    for community_id, community_data in search_index_data.items():
                        self.community_search_index[community_id] = community_data
                else:
                    self.logger.warning(f"Unexpected search index format: {type(search_index_data)}")
            
            self.logger.info(f"Loaded search index for {len(self.community_search_index)} communities")
            
        except Exception as e:
            self.logger.error(f"Failed to load community search index: {e}")
    
    async def _retrieve_hyde_embeddings_enhanced(self) -> Dict[str, Dict[str, Any]]:
        """Retrieve HyDE embeddings and metadata."""
        try:
            # Retrieve community embeddings
            query = """

            MATCH (c:Community)

            WHERE c.hyde_embeddings IS NOT NULL

            OPTIONAL MATCH (meta:CommunitiesMetadata)

            RETURN c.id as community_id,

                   c.hyde_embeddings as hyde_embeddings,

                   c.summary as summary,

                   c.key_entities as key_entities,

                   c.member_ids as member_ids,

                   size(c.hyde_embeddings) as embedding_size,

                   meta.modularity_score as global_modularity_score

            """
            
            results = self.neo4j_conn.execute_query(query)
            hyde_embeddings = {}
            
            for record in results:
                community_id = record['community_id']
                embeddings_data = record.get('hyde_embeddings')
                
                if embeddings_data and community_id:
                    hyde_embeddings[community_id] = {
                        'embeddings': embeddings_data,
                        'summary': record.get('summary', ''),
                        'key_entities': record.get('key_entities', []),
                        'member_ids': record.get('member_ids', []),
                        'embedding_size': record.get('embedding_size', 0),
                        'global_modularity_score': record.get('global_modularity_score', 0.0),
                        'embedding_type': 'hyde'
                    }
            
            self.logger.info(f"Retrieved enhanced HyDE embeddings for {len(hyde_embeddings)} communities")
            return hyde_embeddings
            
        except Exception as e:
            self.logger.error(f"Failed to retrieve enhanced HyDE embeddings: {e}")
            # Retry logic for embeddings
            self.logger.info("Attempting retry for HyDE embeddings...")
            try:
                import time
                time.sleep(2)  # Brief delay before retry
                results = self.neo4j_conn.execute_query(query)
                hyde_embeddings = {}
                
                for record in results:
                    community_id = record['community_id']
                    embeddings_data = record.get('hyde_embeddings')
                    
                    if embeddings_data and community_id:
                        hyde_embeddings[community_id] = {
                            'embeddings': embeddings_data,
                            'summary': record.get('summary', ''),
                            'key_entities': record.get('key_entities', []),
                            'member_ids': record.get('member_ids', []),
                            'embedding_size': record.get('embedding_size', 0),
                            'global_modularity': record.get('global_modularity_score', 0.0)
                        }
                
                self.logger.info(f"Retry successful: Retrieved enhanced HyDE embeddings for {len(hyde_embeddings)} communities")
                return hyde_embeddings
                
            except Exception as retry_error:
                self.logger.error(f"Retry also failed: {retry_error}")
                return {}
    
    def _compute_hyde_similarities_enhanced(self,

                                          query_embedding: List[float],

                                          hyde_embeddings: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, float]]:
        """

        Enhanced similarity computation with global modularity weighting.

        

        Calculates similarity scores between query embedding and community embeddings.

        """
        similarities = {}
        query_vec = np.array(query_embedding)
        query_norm = np.linalg.norm(query_vec)
        
        if query_norm == 0:
            self.logger.warning("Query embedding has zero norm")
            return similarities
        
        for community_id, embedding_data in hyde_embeddings.items():
            embeddings_list = embedding_data['embeddings']
            global_modularity = embedding_data.get('global_modularity_score', 0.0)
            
            max_similarity = 0.0
            
            # Compute similarity
            try:
                # Parse embedding string
                if isinstance(embeddings_list, str):
                    embeddings_list = json.loads(embeddings_list)
                
                # Process embeddings
                if isinstance(embeddings_list, list) and len(embeddings_list) > 0:
                    # Use first embedding
                    hyde_vec = np.array(embeddings_list[0] if isinstance(embeddings_list[0], list) else embeddings_list)
                else:
                    hyde_vec = np.array(embeddings_list)
                
                hyde_norm = np.linalg.norm(hyde_vec)
                
                if hyde_norm > 0:
                    # Calculate similarity
                    base_similarity = np.dot(query_vec, hyde_vec) / (query_norm * hyde_norm)
                    
                    # Apply weighting
                    weighted_similarity = base_similarity * (1 + 0.2 * global_modularity)
                    max_similarity = weighted_similarity
                        
            except Exception as e:
                self.logger.warning(f"Error computing similarity for community {community_id}: {e}")
                continue
            
            similarities[community_id] = {
                'similarity': max_similarity,
                'global_modularity_score': global_modularity,
                'embedding_size': embedding_data.get('embedding_size', 0)
            }
        
        self.logger.info(f"Computed enhanced similarities for {len(similarities)} communities")
        return similarities
    
    def _rank_communities_enhanced(self,

                                 similarities: Dict[str, Dict[str, float]],

                                 routing_result: DriftRoutingResult) -> List[Tuple[str, Dict[str, float]]]:
        """

        Enhanced ranking using global modularity and similarity.

        

        Ranks communities based on a weighted combination of similarity score and modularity.

        """
        
        # Rank primarily by similarity, with modularity as secondary factor
        
        def ranking_score(item):
            _, scores = item
            similarity = scores['similarity']
            global_modularity = scores['global_modularity_score']
            
            # Weighted combination (similarity is primary)
            return 0.8 * similarity + 0.2 * global_modularity
        
        # Sort by combined ranking score
        ranked = sorted(similarities.items(), key=ranking_score, reverse=True)
        
        # Apply similarity threshold
        similarity_threshold = routing_result.parameters.get('similarity_threshold', 0.7)
        filtered_ranked = [
            (cid, scores) for cid, scores in ranked 
            if scores['similarity'] >= similarity_threshold
        ]
        
        self.logger.info(f"Enhanced ranking: {len(filtered_ranked)} communities above threshold {similarity_threshold}")
        return filtered_ranked
    
    def _apply_termination_criteria(self,

                                  ranked_communities: List[Tuple[str, Dict[str, float]]],

                                  routing_result: DriftRoutingResult) -> List[Tuple[str, Dict[str, float]]]:
        """

        Apply termination criteria for community selection.

        

        Limits the number of communities selected based on threshold parameters.

        """
        
        # Get termination criteria from routing or defaults
        max_communities = routing_result.parameters.get('max_communities', 3)
        min_global_modularity = routing_result.parameters.get('min_global_modularity', 0.3)
        
        # Apply criteria
        filtered = []
        for community_id, scores in ranked_communities:
            if len(filtered) >= max_communities:
                break
                
            # Check global modularity threshold
            if scores['global_modularity_score'] >= min_global_modularity:
                filtered.append((community_id, scores))
        
        self.logger.info(f"Applied termination criteria: {len(filtered)} communities selected")
        return filtered
    
    async def _fetch_community_details_enhanced(self,

                                              ranked_communities: List[Tuple[str, Dict[str, float]]]) -> List[CommunityResult]:
        """

        Fetch comprehensive community details with all properties.

        

        Retrieves detailed information about selected communities including summaries,

        key entities, and member IDs.

        """
        community_results = []
        
        for community_id, scores in ranked_communities:
            try:
                # Query the Community node directly by ID (since embedding communities have id=community_id)
                detail_query = """

                MATCH (c:Community)

                WHERE c.id = $community_id AND c.hyde_embeddings IS NOT NULL

                OPTIONAL MATCH (meta:CommunitiesMetadata)

                RETURN c.summary as summary,

                       c.key_entities as key_entities,

                       c.member_ids as member_ids,

                       c.internal_edges as internal_edges,

                       c.density as density,

                       c.avg_degree as avg_degree,

                       c.level as level,

                       meta.modularity_score as modularity_score,

                       CASE WHEN c.member_ids IS NOT NULL THEN size(c.member_ids) ELSE 0 END as member_count,

                       c.id as id

                LIMIT 1

                """
                
                results = self.neo4j_conn.execute_query(
                    detail_query, 
                    {'community_id': community_id}
                )
                
                if results:
                    record = results[0]
                    
                    # Create enhanced community result with actual available data from Neo4j
                    community_result = CommunityResult(
                        community_id=community_id,
                        similarity_score=scores['similarity'],
                        summary=record.get('summary', ''),
                        key_entities=record.get('key_entities', []),
                        member_ids=record.get('member_ids', []),
                        modularity_score=record.get('modularity_score', 0.0),
                        level=record.get('level', 1),
                        internal_edges=record.get('internal_edges', 0),
                        member_count=record.get('member_count', 0),
                        confidence_score=scores.get('confidence_score', 0.5),
                        search_index='',
                        termination_criteria={},
                        centrality_stats={
                            'avg_degree': record.get('avg_degree', 0.0),
                            'density': record.get('density', 0.0)
                        }
                    )
                    
                    community_results.append(community_result)
                    
            except Exception as e:
                self.logger.error(f"Failed to fetch details for community {community_id}: {e}")
                continue
        
        self.logger.info(f"Fetched enhanced details for {len(community_results)} communities")
        return community_results
    
    async def _extract_community_data_enhanced(self,

                                             communities: List[CommunityResult]) -> Dict[str, Any]:
        """

        Step 7: Enhanced data extraction with centrality measures.

        

        Extracts:

        - Entities with degree/betweenness/closeness centrality

        - Relationships with confidence scores

        - Community statistics and properties

        """
        try:
            all_entities = []
            all_relationships = []
            community_stats = []
            
            for community in communities:
                # Extract entities with centrality measures
                entities = await self._extract_entities_with_centrality(community)
                all_entities.extend(entities)
                
                # Extract relationships with properties
                relationships = await self._extract_relationships_enhanced(community)
                all_relationships.extend(relationships)
                
                # Collect community statistics
                community_stats.append({
                    'community_id': community.community_id,
                    'member_count': community.member_count,
                    'modularity_score': community.modularity_score,
                    'confidence_score': community.confidence_score,
                    'centrality_stats': community.centrality_stats
                })
            
            extracted_data = {
                'entities': all_entities,
                'relationships': all_relationships,
                'community_stats': community_stats,
                'extraction_metadata': {
                    'communities_processed': len(communities),
                    'entities_extracted': len(all_entities),
                    'relationships_extracted': len(all_relationships),
                    'timestamp': datetime.now().isoformat()
                }
            }
            
            self.logger.info(f"Enhanced extraction completed: {len(all_entities)} entities, {len(all_relationships)} relationships")
            return extracted_data
            
        except Exception as e:
            self.logger.error(f"Enhanced data extraction failed: {e}")
            return {'entities': [], 'relationships': [], 'community_stats': []}
    
    async def _extract_entities_with_centrality(self,

                                              community: CommunityResult) -> List[EntityResult]:
        """

        Extract entities with comprehensive centrality measures.

        

        Retrieves entities from the community with their associated centrality metrics.

        """
        try:
            # Use member_ids for direct access if available
            member_ids = community.member_ids if community.member_ids else []
            
            if member_ids:
                # Direct member access query based on actual schema
                entity_query = """

                MATCH (n)

                WHERE n.id IN $member_ids

                  AND n.name IS NOT NULL 

                  AND n.content IS NOT NULL

                RETURN n.id as entity_id,

                       n.name as name,

                       n.content as content,

                       n.confidence as confidence,

                       n.degree_centrality as degree_centrality,

                       n.betweenness_centrality as betweenness_centrality,

                       n.closeness_centrality as closeness_centrality,

                       labels(n) as node_types

                ORDER BY n.degree_centrality DESC

                """
                
                results = self.neo4j_conn.execute_query(
                    entity_query,
                    {'member_ids': member_ids}
                )
            else:
                # Fallback: find entities using community_id pattern matching
                entity_query = """

                MATCH (n)

                WHERE n.community_id IS NOT NULL

                  AND n.name IS NOT NULL 

                  AND n.content IS NOT NULL

                RETURN n.id as entity_id,

                       n.name as name,

                       n.content as content,

                       n.confidence as confidence,

                       n.degree_centrality as degree_centrality,

                       n.betweenness_centrality as betweenness_centrality,

                       n.closeness_centrality as closeness_centrality,

                       labels(n) as node_types

                ORDER BY n.degree_centrality DESC

                LIMIT 20

                """
                
                results = self.neo4j_conn.execute_query(entity_query)
            
            entities = []
            for record in results:
                entity = EntityResult(
                    entity_id=record['entity_id'],
                    name=record.get('name', ''),
                    content=record.get('content', ''),
                    confidence=record.get('confidence', 0.0),
                    degree_centrality=record.get('degree_centrality', 0.0),
                    betweenness_centrality=record.get('betweenness_centrality', 0.0),
                    closeness_centrality=record.get('closeness_centrality', 0.0),
                    community_id=community.community_id,
                    node_type=record.get('node_types', ['Unknown'])[0] if record.get('node_types') else 'Unknown'
                )
                entities.append(entity)
            
            return entities
            
        except Exception as e:
            self.logger.error(f"Failed to extract entities for community {community.community_id}: {e}")
            return []
    
    async def _extract_relationships_enhanced(self,

                                            community: CommunityResult) -> List[RelationshipResult]:
        """

        Extract relationships with enhanced properties.

        

        Retrieves relationship data between entities within the specified community.

        """
        try:
            relationship_query = """

            MATCH (a)-[r]->(b)

            WHERE a.community_id = $community_id 

              AND b.community_id = $community_id

              AND r.confidence > 0.5

            RETURN startNode(r).id as start_node,

                   endNode(r).id as end_node,

                   type(r) as relationship_type,

                   r.confidence as confidence

            ORDER BY r.confidence DESC

            LIMIT 50

            """
            
            results = self.neo4j_conn.execute_query(
                relationship_query,
                {'community_id': community.community_id}
            )
            
            relationships = []
            for record in results:
                relationship = RelationshipResult(
                    start_node=record['start_node'],
                    end_node=record['end_node'],
                    relationship_type=record['relationship_type'],
                    confidence=record.get('confidence', 0.0)
                )
                relationships.append(relationship)
            
            return relationships
            
        except Exception as e:
            self.logger.error(f"Failed to extract relationships for community {community.community_id}: {e}")
            return []
    
    async def _generate_initial_answer_enhanced(self, 

                                              extracted_data: Dict[str, Any],

                                              routing_result: DriftRoutingResult) -> Dict[str, Any]:
        """

        Step 8: Context-aware initial answer generation.

        

        Uses:

        - Entity importance from centrality measures

        - Relationship confidence for evidence strength

        - Community statistics for context sizing

        """
        try:
            entities = extracted_data['entities']
            relationships = extracted_data['relationships']
            community_stats = extracted_data['community_stats']
            
            # Rank entities by importance (centrality measures)
            important_entities = sorted(
                entities, 
                key=lambda e: (e.degree_centrality + e.betweenness_centrality) / 2,
                reverse=True
            )[:10]
            
            # Select high-confidence relationships
            strong_relationships = [
                r for r in relationships 
                if r.confidence >= 0.7
            ]
            
            # Prepare context for LLM
            llm_context = self._prepare_llm_context_enhanced(
                important_entities, strong_relationships, community_stats, routing_result
            )
            
            # Generate initial answer using configured LLM
            llm_response = await self._generate_llm_answer(llm_context, routing_result)
            
            initial_answer = {
                'content': llm_response['answer'],
                'llm_context': llm_context,
                'context_used': {
                    'important_entities': len(important_entities),
                    'strong_relationships': len(strong_relationships),
                    'communities_analyzed': len(community_stats)
                },
                'confidence_metrics': {
                    'avg_entity_centrality': np.mean([e.degree_centrality for e in important_entities]) if important_entities else 0,
                    'avg_relationship_confidence': np.mean([r.confidence for r in strong_relationships]) if strong_relationships else 0,
                    'avg_community_modularity': np.mean([c['modularity_score'] for c in community_stats]) if community_stats else 0,
                    'llm_confidence': llm_response['confidence']
                },
                'follow_up_questions': llm_response['follow_up_questions'],
                'reasoning': llm_response['reasoning']
            }
            
            self.logger.info("Enhanced initial answer generated with comprehensive context")
            return initial_answer
            
        except Exception as e:
            self.logger.error(f"Enhanced answer generation failed: {e}")
            return {'content': 'Error generating initial answer', 'error': str(e)}
    
    def _prepare_llm_context_enhanced(self, 

                                    entities: List[EntityResult],

                                    relationships: List[RelationshipResult],

                                    community_stats: List[Dict[str, Any]],

                                    routing_result: DriftRoutingResult) -> str:
        """Prepare enhanced context for LLM with comprehensive information."""
        
        context_parts = [
            f"Query: {routing_result.original_query}",
            f"Search Strategy: {routing_result.search_strategy.value}",
            "",
            "=== IMPORTANT ENTITIES (Use these specific names in your answer) ===",
        ]
        
        for i, entity in enumerate(entities[:10], 1):  # Show more entities
            context_parts.append(
                f"{i}. NAME: '{entity.name}' | Description: {entity.content[:100]}... "
                f"| Centrality: {entity.degree_centrality:.3f} | Confidence: {entity.confidence:.3f}"
            )
        
        context_parts.extend([
            "",
            "=== KEY RELATIONSHIPS (Use these connections in your answer) ===",
        ])
        
        for i, rel in enumerate(relationships[:8], 1):  # Show more relationships
            context_parts.append(
                f"{i}. '{rel.start_node}' --[{rel.relationship_type}]--> '{rel.end_node}' "
                f"| Confidence: {rel.confidence:.3f}"
            )
        
        # Add quick reference list of all entity names
        entity_names = [entity.name for entity in entities[:15]]
        context_parts.extend([
            "",
            "=== ENTITY NAMES FOR REFERENCE ===",
            f"Available entities: {', '.join(entity_names)}",
            "",
            "=== COMMUNITY STATISTICS ===",
        ])
        
        for stat in community_stats:
            context_parts.append(
                f"Community {stat['community_id']}: {stat['member_count']} members, "
                f"modularity: {stat['modularity_score']:.3f}"
            )
        
        context_parts.extend([
            "",
            "REMEMBER: Use the specific entity names listed above in your answer!"
        ])
        
        return "\n".join(context_parts)
    
    async def _generate_llm_answer(self, 

                                 context: str, 

                                 routing_result: DriftRoutingResult) -> Dict[str, Any]:
        """

        Generate actual LLM response using the configured LLM.

        

        Uses the LLM from GraphRAGSetup to generate answers with follow-up questions.

        """
        try:
            # Construct comprehensive prompt for LLM
            prompt = f"""

You are an expert knowledge analyst. Answer the user's query using SPECIFIC NAMES and information from the graph data provided below.



IMPORTANT: Use the actual entity names, organization names, and relationship details from the graph data. Do not give generic answers.



GRAPH DATA CONTEXT:

{context}



INSTRUCTIONS:

1. Answer using SPECIFIC ENTITY NAMES from the "IMPORTANT ENTITIES" section above

2. Reference actual relationships and organizations mentioned in the graph data

3. If the query asks for members/organizations, LIST THE ACTUAL NAMES from the entities

4. Use confidence scores and centrality measures as evidence strength indicators

5. Generate follow-up questions based on the specific entities found



RESPONSE FORMAT:

Answer: [Use specific names and details from the graph data above]

Confidence: [0.0-1.0]

Reasoning: [Why these specific entities answer the query]

Follow-up Questions:

1. [Specific question about entities found]

2. [Question about relationships discovered]

3. [Question about community connections]

4. [Question for deeper exploration]

5. [Question about related entities]

"""

            # Call the configured LLM
            llm_response = await self.setup.llm.acomplete(prompt)
            response_text = llm_response.text
            
            # Parse LLM response
            parsed_response = self._parse_llm_response(response_text)
            
            self.logger.info(f"LLM generated answer with confidence: {parsed_response['confidence']}")
            return parsed_response
            
        except Exception as e:
            self.logger.error(f"LLM answer generation failed: {e}")
            # Fallback response
            return {
                'answer': f"Based on the graph analysis, I found relevant information but encountered an issue generating the full response: {str(e)}",
                'confidence': 0.3,
                'reasoning': "LLM generation encountered an error, providing basic analysis from graph data.",
                'follow_up_questions': [
                    "What specific aspects would you like me to explore further?",
                    "Are there particular entities or relationships of interest?",
                    "Should I focus on a specific community or time period?"
                ]
            }
    
    def _parse_llm_response(self, response_text: str) -> Dict[str, Any]:
        """Parse structured LLM response into components."""
        try:
            lines = response_text.strip().split('\n')
            
            answer = ""
            confidence = 0.5
            reasoning = ""
            follow_up_questions = []
            
            current_section = None
            
            for line in lines:
                line = line.strip()
                
                if line.startswith("Answer:"):
                    current_section = "answer"
                    answer = line.replace("Answer:", "").strip()
                elif line.startswith("Confidence:"):
                    confidence_text = line.replace("Confidence:", "").strip()
                    try:
                        confidence = float(confidence_text)
                    except (ValueError, TypeError):
                        confidence = 0.5
                elif line.startswith("Reasoning:"):
                    current_section = "reasoning"
                    reasoning = line.replace("Reasoning:", "").strip()
                elif line.startswith("Follow-up Questions:"):
                    current_section = "questions"
                elif current_section == "answer" and line:
                    answer += " " + line
                elif current_section == "reasoning" and line:
                    reasoning += " " + line
                elif current_section == "questions" and line.startswith(("1.", "2.", "3.", "4.", "5.")):
                    question = line[2:].strip()  # Remove "1. " etc.
                    follow_up_questions.append(question)
            
            return {
                'answer': answer.strip() if answer else "Unable to generate answer from available context.",
                'confidence': max(0.0, min(1.0, confidence)),  # Clamp between 0-1
                'reasoning': reasoning.strip() if reasoning else "Analysis based on graph structure and entity relationships.",
                'follow_up_questions': follow_up_questions if follow_up_questions else [
                    "What additional information would be helpful?",
                    "Are there specific aspects to explore further?",
                    "Should I analyze different communities or relationships?"
                ]
            }
            
        except Exception as e:
            self.logger.error(f"Failed to parse LLM response: {e}")
            return {
                'answer': response_text[:500] if response_text else "No response generated.",
                'confidence': 0.4,
                'reasoning': "Direct LLM output due to parsing issues.",
                'follow_up_questions': ["What would you like to know more about?"]
            }


# Exports
__all__ = ['CommunitySearchEngine', 'CommunityResult', 'EntityResult', 'RelationshipResult']