| from loguru import logger | |
| class CoTSolver: | |
| def __init__(self, | |
| querier, | |
| system_prompt: str = None): | |
| """ | |
| Initializes the CoTSolver with the given parameters. | |
| Args: | |
| querier: The querier object used for querying the model. | |
| system_prompt (str, optional): The system prompt to be used. Defaults to None. | |
| parse_feedback (bool, optional): Flag to enable parsing feedback. Defaults to False. | |
| check_feedback (bool, optional): Flag to enable checking feedback. Defaults to False. | |
| """ | |
| self.querier = querier | |
| self.cost = 0 | |
| self.detailed_cost = [] | |
| self.system_prompt = system_prompt | |
| def build_query(self, problem): | |
| """ | |
| Constructs an initial query message for the given problem. | |
| Args: | |
| problem (object): The problem instance to build the query for. It should have a method and a | |
| get_formatting_instructions method. | |
| Returns: | |
| list: A list of message dictionaries formatted for the query. The list includes a system prompt if | |
| provided, followed by the user prompt containing the problem description and formatting instructions. | |
| """ | |
| prompt, image_path = str(problem[0]), problem[1] | |
| messages = [] | |
| if self.system_prompt is not None: | |
| messages.append({"role": "system", "content": self.system_prompt}) | |
| messages.append({"role": "user", "content": prompt}) | |
| return messages, image_path | |
| def build_queries(self, problems): | |
| """ | |
| Build a list of queries from a list of problems. | |
| Args: | |
| problems (list): A list of problem instances. | |
| Returns: | |
| list: A list of queries generated from the given problems. | |
| """ | |
| queries = [] | |
| for problem in problems: | |
| queries.append(self.build_query(problem)) | |
| return queries | |
| def add_response(self, query, response): | |
| query, _ = query | |
| if isinstance(response, tuple) and response[0] is None: | |
| query.append({"role": "api_error", "content": str(response[1])}) | |
| else: | |
| query.append({"role": "assistant", "content": response}) | |
| return query | |
| def solve(self, problems): | |
| """ | |
| Solves the initial round of problems by building queries, running them, and appending responses. | |
| Args: | |
| problems (list): A list of problems to be solved. | |
| Returns: | |
| list: A list of queries with appended responses from the assistant. | |
| """ | |
| logger.info("Solving problems.") | |
| queries = self.build_queries(problems) | |
| self.cost = 0 | |
| self.detailed_cost = [{ | |
| "cost": 0, | |
| "input_tokens": 0, | |
| "output_tokens": 0, | |
| } for _ in range(len(problems))] | |
| for idx, response, detailed_cost in self.querier.run_queries(queries): | |
| messages = self.add_response(queries[idx], response) | |
| self.detailed_cost[idx]["cost"] += detailed_cost["cost"] | |
| self.detailed_cost[idx]["input_tokens"] += detailed_cost["input_tokens"] | |
| self.detailed_cost[idx]["output_tokens"] += detailed_cost["output_tokens"] | |
| yield idx, messages, detailed_cost | |