Coverage for tinytroupe / extraction / results_reducer.py: 0%

1import pandas as pd 

2 

3from tinytroupe.extraction import logger 

4from tinytroupe.agent import TinyPerson 

5 

6 

7class ResultsReducer: 

8 

9 def __init__(self): 

10 self.results = {} 

11 

12 self.rules = {} 

13 

14 def add_reduction_rule(self, trigger: str, func: callable): 

15 if trigger in self.rules: 

16 raise Exception(f"Rule for {trigger} already exists.") 

17 

18 self.rules[trigger] = func 

19 

20 def reduce_agent(self, agent: TinyPerson) -> list: 

21 reduction = [] 

22 for message in agent.episodic_memory.retrieve_all(): 

23 if message['role'] == 'system': 

24 continue # doing nothing for `system` role yet at least 

25 

26 elif message['role'] == 'user': 

27 # User role is related to stimuli only 

28 stimulus_type = message['content']['stimuli'][0].get('type', None) 

29 stimulus_content = message['content']['stimuli'][0].get('content', None) 

30 stimulus_source = message['content']['stimuli'][0].get('source', None) 

31 stimulus_timestamp = message['simulation_timestamp'] 

32 

33 if stimulus_type in self.rules: 

34 extracted = self.rules[stimulus_type](focus_agent=agent, source_agent=TinyPerson.get_agent_by_name(stimulus_source), target_agent=agent, kind='stimulus', event=stimulus_type, content=stimulus_content, timestamp=stimulus_timestamp) 

35 if extracted is not None: 

36 reduction.append(extracted) 

37 

38 elif message['role'] == 'assistant': 

39 # Assistant role is related to actions only 

40 if 'action' in message['content']: 

41 action_type = message['content']['action'].get('type', None) 

42 action_content = message['content']['action'].get('content', None) 

43 action_target = message['content']['action'].get('target', None) 

44 action_timestamp = message['simulation_timestamp'] 

45 

46 if action_type in self.rules: 

47 extracted = self.rules[action_type](focus_agent=agent, source_agent=agent, target_agent=TinyPerson.get_agent_by_name(action_target), kind='action', event=action_type, content=action_content, timestamp=action_timestamp) 

48 if extracted is not None: 

49 reduction.append(extracted) 

50 

51 return reduction 

52 

53 def reduce_agent_to_dataframe(self, agent: TinyPerson, column_names: list=None) -> pd.DataFrame: 

54 reduction = self.reduce_agent(agent) 

55 return pd.DataFrame(reduction, columns=column_names)