Update app.py

app.py

@@ -47,6 +47,7 @@ class RelationshipEngine:
         self.setup_database()
 
     def setup_database(self):
+        # Events table
         self.conn.execute('''
             CREATE TABLE IF NOT EXISTS events (
                 id INTEGER PRIMARY KEY,
@@ -56,46 +57,171 @@ class RelationshipEngine:
             )
         ''')
         
+        # Entities table
         self.conn.execute('''
-            CREATE TABLE IF NOT EXISTS relationships (
+            CREATE TABLE IF NOT EXISTS entities (
                 id INTEGER PRIMARY KEY,
-                source_event_id INTEGER,
-                target_event_id INTEGER,
+                entity_text TEXT,
+                entity_type TEXT,    -- person, organization, location, hashtag, temporal
+                first_seen DATETIME,
+                last_seen DATETIME,
+                frequency INTEGER DEFAULT 1,
+                confidence REAL
+            )
+        ''')
+        
+        # Event-Entity relationships
+        self.conn.execute('''
+            CREATE TABLE IF NOT EXISTS event_entities (
+                event_id INTEGER,
+                entity_id INTEGER,
+                FOREIGN KEY (event_id) REFERENCES events(id),
+                FOREIGN KEY (entity_id) REFERENCES entities(id),
+                PRIMARY KEY (event_id, entity_id)
+            )
+        ''')
+        
+        # Entity relationships (e.g., person-organization affiliations)
+        self.conn.execute('''
+            CREATE TABLE IF NOT EXISTS entity_relationships (
+                id INTEGER PRIMARY KEY,
+                source_entity_id INTEGER,
+                target_entity_id INTEGER,
                 relationship_type TEXT,
                 confidence REAL,
-                FOREIGN KEY (source_event_id) REFERENCES events(id),
-                FOREIGN KEY (target_event_id) REFERENCES events(id)
+                first_seen DATETIME,
+                last_seen DATETIME,
+                FOREIGN KEY (source_entity_id) REFERENCES entities(id),
+                FOREIGN KEY (target_entity_id) REFERENCES entities(id)
             )
         ''')
+        
+        self.conn.commit()
+
+    def store_entities(self, event_id, entities_dict):
+        now = datetime.now().isoformat()
+        
+        for entity_type, entities in entities_dict.items():
+            if not isinstance(entities, list):
+                continue
+                
+            for entity_text in entities:
+                # Check if entity exists
+                cursor = self.conn.execute(
+                    'SELECT id, frequency FROM entities WHERE entity_text = ? AND entity_type = ?',
+                    (entity_text, entity_type)
+                )
+                result = cursor.fetchone()
+                
+                if result:
+                    # Update existing entity
+                    entity_id, freq = result
+                    self.conn.execute('''
+                        UPDATE entities 
+                        SET frequency = ?, last_seen = ?
+                        WHERE id = ?
+                    ''', (freq + 1, now, entity_id))
+                else:
+                    # Insert new entity
+                    cursor = self.conn.execute('''
+                        INSERT INTO entities (entity_text, entity_type, first_seen, last_seen, confidence)
+                        VALUES (?, ?, ?, ?, ?)
+                    ''', (entity_text, entity_type, now, now, 1.0))
+                    entity_id = cursor.lastrowid
+                
+                # Create event-entity relationship
+                self.conn.execute('''
+                    INSERT OR IGNORE INTO event_entities (event_id, entity_id)
+                    VALUES (?, ?)
+                ''', (event_id, entity_id))
+        
         self.conn.commit()
 
     def find_related_events(self, event_data):
-        # Find events with similar entities
-        cursor = self.conn.execute('''
-            SELECT * FROM events 
-            WHERE text LIKE ? 
-            ORDER BY timestamp DESC 
+        # Find events sharing entities
+        entity_texts = []
+        for entity_type, entities in event_data.get('entities', {}).items():
+            if isinstance(entities, list):
+                entity_texts.extend(entities)
+
+        if not entity_texts:
+            return []
+
+        # Build query using entity relationships
+        query = '''
+            SELECT DISTINCT e.*, COUNT(ee.entity_id) as shared_entities
+            FROM events e
+            JOIN event_entities ee ON e.id = ee.event_id
+            JOIN entities ent ON ee.entity_id = ent.id
+            WHERE ent.entity_text IN ({})
+            GROUP BY e.id
+            ORDER BY shared_entities DESC, e.timestamp DESC
             LIMIT 5
-        ''', (f"%{event_data.get('text', '')}%",))
-        
-        related_events = cursor.fetchall()
-        return related_events
+        '''.format(','.join('?' * len(entity_texts)))
+
+        cursor = self.conn.execute(query, entity_texts)
+        return cursor.fetchall()
+
+    def find_entity_relationships(self, entity_id):
+        # Find direct relationships
+        query = '''
+            SELECT 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 entity_relationships er
+            JOIN entities e1 ON er.source_entity_id = e1.id
+            JOIN entities e2 ON er.target_entity_id = e2.id
+            WHERE er.source_entity_id = ? OR er.target_entity_id = ?
+        '''
+        cursor = self.conn.execute(query, (entity_id, entity_id))
+        return cursor.fetchall()
 
-    def calculate_relationship_confidence(self, event1, event2):
-        # Simple similarity-based confidence
-        base_confidence = 0.0
+    def update_entity_relationships(self, event_id):
+        # Find all entities in the event
+        query = '''
+            SELECT e.id, e.entity_text, e.entity_type
+            FROM entities e
+            JOIN event_entities ee ON e.id = ee.entity_id
+            WHERE ee.event_id = ?
+        '''
+        cursor = self.conn.execute(query, (event_id,))
+        entities = cursor.fetchall()
         
-        # Entity overlap increases confidence
-        if set(event1.get('entities', {}).get('people', [])) & set(event2.get('entities', {}).get('people', [])):
-            base_confidence += 0.3
-            
-        if set(event1.get('entities', {}).get('organizations', [])) & set(event2.get('entities', {}).get('organizations', [])):
-            base_confidence += 0.3
-            
-        if set(event1.get('entities', {}).get('locations', [])) & set(event2.get('entities', {}).get('locations', [])):
-            base_confidence += 0.4
-            
-        return min(base_confidence, 1.0)
+        now = datetime.now().isoformat()
+        
+        # Create/update relationships between entities in same event
+        for i, entity1 in enumerate(entities):
+            for entity2 in entities[i+1:]:
+                # Skip same entity type relationships
+                if entity1[2] == entity2[2]:
+                    continue
+                    
+                relationship_type = f"{entity1[2]}_to_{entity2[2]}"
+                
+                # Check if relationship exists
+                cursor = self.conn.execute('''
+                    SELECT id FROM entity_relationships
+                    WHERE (source_entity_id = ? AND target_entity_id = ?)
+                    OR (source_entity_id = ? AND target_entity_id = ?)
+                ''', (entity1[0], entity2[0], entity2[0], entity1[0]))
+                
+                result = cursor.fetchone()
+                if result:
+                    # Update existing relationship
+                    self.conn.execute('''
+                        UPDATE entity_relationships
+                        SET last_seen = ?, confidence = confidence + 0.1
+                        WHERE id = ?
+                    ''', (now, result[0]))
+                else:
+                    # Create new relationship
+                    self.conn.execute('''
+                        INSERT INTO entity_relationships
+                        (source_entity_id, target_entity_id, relationship_type, confidence, first_seen, last_seen)
+                        VALUES (?, ?, ?, ?, ?, ?)
+                    ''', (entity1[0], entity2[0], relationship_type, 0.5, now, now))
+        
+        self.conn.commit()
 
 class EventAnalyzer:
     def __init__(self):