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

87 statements  

« prev     ^ index     » next       coverage.py v7.13.4, created at 2026-02-28 17:48 +0000

1import os 

2import json 

3import chevron 

4import pandas as pd 

5from typing import Union, List 

6 

7from tinytroupe.extraction import logger 

8from tinytroupe.agent import TinyPerson 

9from tinytroupe.environment import TinyWorld 

10 

11from tinytroupe import openai_utils 

12import tinytroupe.utils as utils 

13 

14 

15class ResultsExtractor: 

16 

17 def __init__(self, 

18 extraction_prompt_template_path:str = os.path.join(os.path.dirname(__file__), './prompts/interaction_results_extractor.mustache'), 

19 extraction_objective:str = "The main points present in the agents' interactions history.", 

20 situation:str = "", 

21 fields:List[str] = None, 

22 fields_hints:dict = None, 

23 verbose:bool = False): 

24 """ 

25 Initializes the ResultsExtractor with default parameters. 

26 

27 Args: 

28 extraction_prompt_template_path (str): The path to the extraction prompt template. 

29 extraction_objective (str): The default extraction objective. 

30 situation (str): The default situation to consider. 

31 fields (List[str], optional): The default fields to extract. Defaults to None. 

32 fields_hints (dict, optional): The default hints for the fields to extract. Defaults to None. 

33 verbose (bool, optional): Whether to print debug messages by default. Defaults to False. 

34 """ 

35 self._extraction_prompt_template_path = extraction_prompt_template_path 

36 

37 # Default parameters 

38 self.default_extraction_objective = extraction_objective 

39 self.default_situation = situation 

40 self.default_fields = fields 

41 self.default_fields_hints = fields_hints 

42 self.default_verbose = verbose 

43 

44 # Cache for the last extraction results 

45 self.agent_extraction = {} 

46 self.world_extraction = {} 

47 

48 def extract_results_from_agents(self, 

49 agents:List[TinyPerson], 

50 extraction_objective:str=None, 

51 situation:str =None, 

52 fields:list=None, 

53 fields_hints:dict=None, 

54 verbose:bool=None): 

55 """ 

56 Extracts results from a list of TinyPerson instances. 

57 

58 Args: 

59 agents (List[TinyPerson]): The list of TinyPerson instances to extract results from. 

60 extraction_objective (str): The extraction objective. 

61 situation (str): The situation to consider. 

62 fields (list, optional): The fields to extract. If None, the extractor will decide what names to use.  

63 Defaults to None. 

64 fields_hints (dict, optional): Hints for the fields to extract. Maps field names to strings with the hints. Defaults to None. 

65 verbose (bool, optional): Whether to print debug messages. Defaults to False. 

66 

67  

68 """ 

69 results = [] 

70 for agent in agents: 

71 result = self.extract_results_from_agent(agent, extraction_objective, situation, fields, fields_hints, verbose) 

72 results.append(result) 

73 

74 return results 

75 

76 def extract_results_from_agent(self, 

77 tinyperson:TinyPerson, 

78 extraction_objective:str="The main points present in the agent's interactions history.", 

79 situation:str = "", 

80 fields:list=None, 

81 fields_hints:dict=None, 

82 verbose:bool=None): 

83 """ 

84 Extracts results from a TinyPerson instance. 

85 

86 Args: 

87 tinyperson (TinyPerson): The TinyPerson instance to extract results from. 

88 extraction_objective (str): The extraction objective. 

89 situation (str): The situation to consider. 

90 fields (list, optional): The fields to extract. If None, the extractor will decide what names to use.  

91 Defaults to None. 

92 fields_hints (dict, optional): Hints for the fields to extract. Maps field names to strings with the hints. Defaults to None. 

93 verbose (bool, optional): Whether to print debug messages. Defaults to False. 

94 """ 

95 

96 extraction_objective, situation, fields, fields_hints, verbose = self._get_default_values_if_necessary( 

97 extraction_objective, situation, fields, fields_hints, verbose 

98 ) 

99 

100 messages = [] 

101 

102 rendering_configs = {} 

103 if fields is not None: 

104 rendering_configs["fields"] = ", ".join(fields) 

105 

106 if fields_hints is not None: 

107 rendering_configs["fields_hints"] = list(fields_hints.items()) 

108 

109 messages.append({"role": "system", 

110 "content": chevron.render( 

111 open(self._extraction_prompt_template_path, 'r', encoding='utf-8', errors='replace').read(), 

112 rendering_configs)}) 

113 

114 

115 interaction_history = tinyperson.pretty_current_interactions(max_content_length=None) 

116 

117 extraction_request_prompt = \ 

118f""" 

119## Extraction objective 

120 

121{extraction_objective} 

122 

123## Situation 

124You are considering a single agent, named {tinyperson.name}. Your objective thus refers to this agent specifically. 

125{situation} 

126 

127## Agent Interactions History 

128 

129You will consider an agent's history of interactions, which include stimuli it received as well as actions it  

130performed. 

131 

132{interaction_history} 

133""" 

134 messages.append({"role": "user", "content": extraction_request_prompt}) 

135 

136 next_message = openai_utils.client().send_message(messages, temperature=0.0, frequency_penalty=0.0, presence_penalty=0.0) 

137 

138 debug_msg = f"Extraction raw result message: {next_message}" 

139 logger.debug(debug_msg) 

140 if verbose: 

141 print(debug_msg) 

142 

143 if next_message is not None: 

144 result = utils.extract_json(next_message["content"]) 

145 else: 

146 result = None 

147 

148 # cache the result 

149 self.agent_extraction[tinyperson.name] = result 

150 

151 return result 

152 

153 

154 def extract_results_from_world(self, 

155 tinyworld:TinyWorld, 

156 extraction_objective:str="The main points that can be derived from the agents conversations and actions.", 

157 situation:str="", 

158 fields:list=None, 

159 fields_hints:dict=None, 

160 verbose:bool=None): 

161 """ 

162 Extracts results from a TinyWorld instance. 

163 

164 Args: 

165 tinyworld (TinyWorld): The TinyWorld instance to extract results from. 

166 extraction_objective (str): The extraction objective. 

167 situation (str): The situation to consider. 

168 fields (list, optional): The fields to extract. If None, the extractor will decide what names to use.  

169 Defaults to None. 

170 verbose (bool, optional): Whether to print debug messages. Defaults to False. 

171 """ 

172 

173 extraction_objective, situation, fields, fields_hints, verbose = self._get_default_values_if_necessary( 

174 extraction_objective, situation, fields, fields_hints, verbose 

175 ) 

176 

177 messages = [] 

178 

179 rendering_configs = {} 

180 if fields is not None: 

181 rendering_configs["fields"] = ", ".join(fields) 

182 

183 if fields_hints is not None: 

184 rendering_configs["fields_hints"] = list(fields_hints.items()) 

185 

186 messages.append({"role": "system", 

187 "content": chevron.render( 

188 open(self._extraction_prompt_template_path, 'r', encoding='utf-8', errors='replace').read(), 

189 rendering_configs)}) 

190 

191 # TODO: either summarize first or break up into multiple tasks 

192 interaction_history = tinyworld.pretty_current_interactions(max_content_length=None) 

193 

194 extraction_request_prompt = \ 

195f""" 

196## Extraction objective 

197 

198{extraction_objective} 

199 

200## Situation 

201You are considering various agents. 

202{situation} 

203 

204## Agents Interactions History 

205 

206You will consider the history of interactions from various agents that exist in an environment called {tinyworld.name}.  

207Each interaction history includes stimuli the corresponding agent received as well as actions it performed. 

208 

209{interaction_history} 

210""" 

211 messages.append({"role": "user", "content": extraction_request_prompt}) 

212 

213 next_message = openai_utils.client().send_message(messages, temperature=0.0) 

214 

215 debug_msg = f"Extraction raw result message: {next_message}" 

216 logger.debug(debug_msg) 

217 if verbose: 

218 print(debug_msg) 

219 

220 if next_message is not None: 

221 result = utils.extract_json(next_message["content"]) 

222 else: 

223 result = None 

224 

225 # cache the result 

226 self.world_extraction[tinyworld.name] = result 

227 

228 return result 

229 

230 def save_as_json(self, filename:str, verbose:bool=False): 

231 """ 

232 Saves the last extraction results as JSON. 

233 

234 Args: 

235 filename (str): The filename to save the JSON to. 

236 verbose (bool, optional): Whether to print debug messages. Defaults to False. 

237 """ 

238 with open(filename, 'w', encoding="utf-8", errors="replace") as f: 

239 json.dump({"agent_extractions": self.agent_extraction, 

240 "world_extraction": self.world_extraction}, f, indent=4) 

241 

242 if verbose: 

243 print(f"Saved extraction results to {filename}") 

244 

245 def _get_default_values_if_necessary(self, 

246 extraction_objective:str, 

247 situation:str, 

248 fields:List[str], 

249 fields_hints:dict, 

250 verbose:bool): 

251 

252 if extraction_objective is None: 

253 extraction_objective = self.default_extraction_objective 

254 

255 if situation is None: 

256 situation = self.default_situation 

257 

258 if fields is None: 

259 fields = self.default_fields 

260 

261 if fields_hints is None: 

262 fields_hints = self.default_fields_hints 

263 

264 if verbose is None: 

265 verbose = self.default_verbose 

266 

267 return extraction_objective, situation, fields, fields_hints, verbose 

268