tinytroupe/ml_models.py

from typing import List, Dict, Any, Optional
import numpy as np
from tinytroupe.agent import TinyPerson
from tinytroupe.agent.social_types import Content, Reaction
from tinytroupe.social_network import NetworkTopology
from tinytroupe.features import ContentFeatureExtractor, PersonaFeatureExtractor, InteractionFeatureExtractor
from tinytroupe.agent.agent_traits import TraitBasedBehaviorModel

class EngagementPredictor:
    """Predicts whether persona will engage with content"""
    
    def __init__(self):
        self.content_extractor = ContentFeatureExtractor()
        self.persona_extractor = PersonaFeatureExtractor()
        self.interaction_extractor = InteractionFeatureExtractor()
        self.trait_model = TraitBasedBehaviorModel()
    
    def predict(self, persona: TinyPerson, content: Content, network: NetworkTopology) -> float:
        """Predict engagement probability"""
        content_features = self.content_extractor.extract(content)
        persona_features = self.persona_extractor.extract(persona)
        interaction_features = self.interaction_extractor.extract(persona, content, network)
        
        # Get base probability from trait-based model
        trait_prob = self.trait_model.compute_action_probability(persona, "engage", content)
        
        # Combine with other signals
        prob = (trait_prob * 0.5 + 
                interaction_features["topic_alignment"] * 0.3 + 
                persona_features["engagement_rate"] * 0.1 +
                content_features["sentiment_score"] * 0.1)
        
        return max(0.0, min(1.0, prob))

class ViralityPredictor:
    def predict_cascade_size(self, content: Content, seed_personas: List[str], network: NetworkTopology) -> int:
        return len(seed_personas) * 2 # Placeholder