update event analyzer

app.py

@@ -259,18 +259,63 @@ class EventAnalyzer:
         return self.ontology.validate_pattern(text, 'temporal')
 
     async def extract_locations(self, text):
-        # First await the entities result
         entities = await self.extract_entities(text)
         ml_locations = entities.get('locations', [])
-        # Get pattern-based locations
         pattern_locations = self.ontology.validate_pattern(text, 'location')
         return list(set(ml_locations + pattern_locations))
 
+    def calculate_confidence(self, entities, temporal_data, related_events):
+        # Base confidence from entity presence
+        base_confidence = min(1.0, (
+            0.2 * bool(entities["people"]) +
+            0.2 * bool(entities["organizations"]) +
+            0.3 * bool(entities["locations"]) +
+            0.3 * bool(temporal_data)
+        ))
+
+        # Adjust confidence based on entity frequency
+        cursor = self.relationship_engine.conn.execute('''
+            SELECT AVG(frequency) as avg_freq 
+            FROM entities 
+            WHERE entity_text IN (
+                SELECT DISTINCT entity_text 
+                FROM entities 
+                WHERE entity_text IN ({})
+            )
+        '''.format(','.join('?' * len([
+            *entities["people"], 
+            *entities["organizations"], 
+            *entities["locations"]
+        ]))))
+        
+        avg_frequency = cursor.fetchone()[0] or 1
+        frequency_boost = min(0.2, (avg_frequency - 1) * 0.05)  # Max 0.2 boost for frequency
+
+        # Adjust confidence based on relationships
+        relationship_confidence = 0
+        if related_events:
+            relationship_scores = []
+            for event in related_events:
+                cursor = self.relationship_engine.conn.execute('''
+                    SELECT COUNT(*) as shared_entities
+                    FROM event_entities ee1
+                    JOIN event_entities ee2 ON ee1.entity_id = ee2.entity_id
+                    WHERE ee1.event_id = ? AND ee2.event_id = ?
+                ''', (event[0], event[0]))  # event[0] is the event_id
+                shared_count = cursor.fetchone()[0]
+                relationship_scores.append(min(0.3, shared_count * 0.1))  # Max 0.3 boost per relationship
+            
+            if relationship_scores:
+                relationship_confidence = max(relationship_scores)
+
+        final_confidence = min(1.0, base_confidence + frequency_boost + relationship_confidence)
+        return final_confidence
+
     async def analyze_event(self, text):
         try:
             # Parallel extraction
             entities_future = self.extract_entities(text)
-            temporal_data = self.extract_temporal(text)  # This is synchronous now
+            temporal_data = self.extract_temporal(text)
             locations_future = self.extract_locations(text)
             
             # Gather async results
@@ -282,30 +327,52 @@ class EventAnalyzer:
             entities['locations'] = locations
             entities['temporal'] = temporal_data
             
-            # Calculate initial confidence
-            confidence = min(1.0, (
-                0.2 * bool(entities["people"]) +
-                0.2 * bool(entities["organizations"]) +
-                0.3 * bool(entities["locations"]) +
-                0.3 * bool(temporal_data)
-            ))
-            
             # Find related events
             related_events = self.relationship_engine.find_related_events({
                 'text': text,
                 'entities': entities
             })
             
-            # Adjust confidence based on relationships
-            if related_events:
-                relationship_confidence = max(
-                    self.relationship_engine.calculate_relationship_confidence(
-                        {'entities': entities}, 
-                        {'text': event[1]}
-                    )
-                    for event in related_events
+            # Calculate confidence with enhanced logic
+            confidence = self.calculate_confidence(entities, temporal_data, related_events)
+            
+            # Store event if confidence meets threshold
+            cursor = None
+            if confidence >= 0.6:
+                cursor = self.relationship_engine.conn.execute(
+                    'INSERT INTO events (text, timestamp, confidence) VALUES (?, ?, ?)',
+                    (text, datetime.now().isoformat(), confidence)
                 )
-                confidence = (confidence + relationship_confidence) / 2
+                event_id = cursor.lastrowid
+                
+                # Store entities and their relationships
+                self.relationship_engine.store_entities(event_id, {
+                    'person': entities['people'],
+                    'organization': entities['organizations'],
+                    'location': entities['locations'],
+                    'temporal': temporal_data,
+                    'hashtag': entities['hashtags']
+                })
+                
+                # Update entity relationships
+                self.relationship_engine.update_entity_relationships(event_id)
+                
+                self.relationship_engine.conn.commit()
+            
+            # Get entity relationships for rich output
+            entity_relationships = []
+            if cursor and cursor.lastrowid:
+                query = '''
+                    SELECT DISTINCT er.*, 
+                           e1.entity_text as source_text, e1.entity_type as source_type,
+                           e2.entity_text as target_text, e2.entity_type as target_type
+                    FROM event_entities ee
+                    JOIN entity_relationships er ON ee.entity_id IN (er.source_entity_id, er.target_entity_id)
+                    JOIN entities e1 ON er.source_entity_id = e1.id
+                    JOIN entities e2 ON er.target_entity_id = e2.id
+                    WHERE ee.event_id = ?
+                '''
+                entity_relationships = self.relationship_engine.conn.execute(query, (cursor.lastrowid,)).fetchall()
             
             result = {
                 "text": text,
@@ -316,25 +383,58 @@ class EventAnalyzer:
                     {
                         "text": event[1],
                         "timestamp": event[2],
-                        "confidence": event[3]
+                        "confidence": event[3],
+                        "shared_entities": event[4] if len(event) > 4 else None
                     }
                     for event in related_events
-                ]
+                ],
+                "entity_relationships": [
+                    {
+                        "type": rel[3],
+                        "source": rel[6],
+                        "target": rel[8],
+                        "confidence": rel[4]
+                    }
+                    for rel in entity_relationships
+                ] if entity_relationships else []
             }
             
-            # Store event if confidence is sufficient
-            if confidence >= 0.6:
-                self.relationship_engine.conn.execute(
-                    'INSERT INTO events (text, timestamp, confidence) VALUES (?, ?, ?)',
-                    (text, datetime.now().isoformat(), confidence)
-                )
-                self.relationship_engine.conn.commit()
-            
             return result
             
         except Exception as e:
             return {"error": str(e)}
 
+    def get_entity_statistics(self):
+        """Get statistics about stored entities and relationships"""
+        stats = {}
+        
+        # Entity counts by type
+        cursor = self.relationship_engine.conn.execute('''
+            SELECT entity_type, COUNT(*) as count, AVG(frequency) as avg_frequency
+            FROM entities
+            GROUP BY entity_type
+        ''')
+        stats['entity_counts'] = cursor.fetchall()
+        
+        # Most frequent entities
+        cursor = self.relationship_engine.conn.execute('''
+            SELECT entity_text, entity_type, frequency
+            FROM entities
+            ORDER BY frequency DESC
+            LIMIT 10
+        ''')
+        stats['frequent_entities'] = cursor.fetchall()
+        
+        # Relationship statistics
+        cursor = self.relationship_engine.conn.execute('''
+            SELECT relationship_type, COUNT(*) as count, AVG(confidence) as avg_confidence
+            FROM entity_relationships
+            GROUP BY relationship_type
+        ''')
+        stats['relationship_stats'] = cursor.fetchall()
+        
+        return stats
+
 # Initialize analyzer
 analyzer = EventAnalyzer()