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
« 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
7from tinytroupe.extraction import logger
8from tinytroupe.agent import TinyPerson
9from tinytroupe.environment import TinyWorld
11from tinytroupe import openai_utils
12import tinytroupe.utils as utils
15class ResultsExtractor:
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.
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
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
44 # Cache for the last extraction results
45 self.agent_extraction = {}
46 self.world_extraction = {}
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.
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.
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)
74 return results
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.
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 """
96 extraction_objective, situation, fields, fields_hints, verbose = self._get_default_values_if_necessary(
97 extraction_objective, situation, fields, fields_hints, verbose
98 )
100 messages = []
102 rendering_configs = {}
103 if fields is not None:
104 rendering_configs["fields"] = ", ".join(fields)
106 if fields_hints is not None:
107 rendering_configs["fields_hints"] = list(fields_hints.items())
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)})
115 interaction_history = tinyperson.pretty_current_interactions(max_content_length=None)
117 extraction_request_prompt = \
118f"""
119## Extraction objective
121{extraction_objective}
123## Situation
124You are considering a single agent, named {tinyperson.name}. Your objective thus refers to this agent specifically.
125{situation}
127## Agent Interactions History
129You will consider an agent's history of interactions, which include stimuli it received as well as actions it
130performed.
132{interaction_history}
133"""
134 messages.append({"role": "user", "content": extraction_request_prompt})
136 next_message = openai_utils.client().send_message(messages, temperature=0.0, frequency_penalty=0.0, presence_penalty=0.0)
138 debug_msg = f"Extraction raw result message: {next_message}"
139 logger.debug(debug_msg)
140 if verbose:
141 print(debug_msg)
143 if next_message is not None:
144 result = utils.extract_json(next_message["content"])
145 else:
146 result = None
148 # cache the result
149 self.agent_extraction[tinyperson.name] = result
151 return result
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.
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 """
173 extraction_objective, situation, fields, fields_hints, verbose = self._get_default_values_if_necessary(
174 extraction_objective, situation, fields, fields_hints, verbose
175 )
177 messages = []
179 rendering_configs = {}
180 if fields is not None:
181 rendering_configs["fields"] = ", ".join(fields)
183 if fields_hints is not None:
184 rendering_configs["fields_hints"] = list(fields_hints.items())
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)})
191 # TODO: either summarize first or break up into multiple tasks
192 interaction_history = tinyworld.pretty_current_interactions(max_content_length=None)
194 extraction_request_prompt = \
195f"""
196## Extraction objective
198{extraction_objective}
200## Situation
201You are considering various agents.
202{situation}
204## Agents Interactions History
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.
209{interaction_history}
210"""
211 messages.append({"role": "user", "content": extraction_request_prompt})
213 next_message = openai_utils.client().send_message(messages, temperature=0.0)
215 debug_msg = f"Extraction raw result message: {next_message}"
216 logger.debug(debug_msg)
217 if verbose:
218 print(debug_msg)
220 if next_message is not None:
221 result = utils.extract_json(next_message["content"])
222 else:
223 result = None
225 # cache the result
226 self.world_extraction[tinyworld.name] = result
228 return result
230 def save_as_json(self, filename:str, verbose:bool=False):
231 """
232 Saves the last extraction results as JSON.
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)
242 if verbose:
243 print(f"Saved extraction results to {filename}")
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):
252 if extraction_objective is None:
253 extraction_objective = self.default_extraction_objective
255 if situation is None:
256 situation = self.default_situation
258 if fields is None:
259 fields = self.default_fields
261 if fields_hints is None:
262 fields_hints = self.default_fields_hints
264 if verbose is None:
265 verbose = self.default_verbose
267 return extraction_objective, situation, fields, fields_hints, verbose