CoderLakshman/Final_Assignment_Template

1from settings import Settings2from models import Question, QuestionAnswerPair3from agent import ManagerAgent4import pandas as pd5import logging6import json7import asyncio8import nest_asyncio9nest_asyncio.apply()10logger = logging.getLogger(__name__)11 12class Runner():13    def __init__(self, settings: Settings):14        self.settings = settings15 16    def _save_pairs(self, pairs: list[QuestionAnswerPair], username: str):17        """Write the question answer pairs to a user-specific file."""18        answers = [pair.model_dump() for pair in pairs if pair is not None]19        file_name = f"answers_{username}.json"20        with open(file_name, "w") as f:21            json.dump(answers, f, indent=4)22 23    def _enrich_question_text(self, item):24        task_id = item.task_id25        file_name = item.file_name26        question_text = (27            f"{item.question} "28            "Think hard to answer. Parse all statements in the question to make a plan. "29            "Your final answer should be a number or as few words as possible. "30            "Only use abbreviations when the question calls for abbreviations. "31            "If needed, use a comma separated list of values; the comma is always followed by a space. "32            f"Critically review your answer before making it the final answer. "33            f"Double check the answer to make sure it meets all format requirements stated in the question. "34            f"task_id: {task_id}."35        )36        if file_name:37            question_text = f"{question_text} file_name: {file_name} (use tools to fetch the file)"38        return question_text39 40    async def _run_agent_async(self, item: Question):41        """Runs the agent asynchronously."""42        task_id = item.task_id43        question_text = self._enrich_question_text(item)44        try:45            answer = await asyncio.to_thread(ManagerAgent(self.settings), question_text)46        except Exception as e:47            logger.error(f"Error running agent on task {task_id}: {e}")48            answer = f"AGENT ERROR: {e}"49        return QuestionAnswerPair(task_id=task_id,50                                  question=item.question, answer=str(answer))51 52    def _assign_questions(self, questions: list[Question]):53        """Runs the asynchronous loop and returns task outputs."""54        tasks = [self._run_agent_async(item) for item in questions]55        return asyncio.gather(*tasks)56 57    def run_agent(self, questions: list[Question], username: str) -> pd.DataFrame:58        """Run the agent(s) async, save answers and return a dataframe"""59        # Assign questions to agents and wait60        try:61            loop = asyncio.get_running_loop()62        except RuntimeError:  # No running loop, create one63            loop = asyncio.new_event_loop()64            asyncio.set_event_loop(loop)65 66        def run_tasks_in_thread():67            question_answer_pairs = loop.run_until_complete(68                self._assign_questions(questions))69            return question_answer_pairs70 71        pairs = run_tasks_in_thread()72 73        # save json to disk and return a dataframe74        self._save_pairs(pairs, username)75        results_log = [pair.model_dump() for pair in pairs if pair is not None]76        if not results_log:77            logger.warning("Agent did not produce any answers to submit.")78 79        return pd.DataFrame(results_log)80