payload.py

"""Runner for Agents to output answers"""

import pandas as pd

from agents.basic_agent import BasicAgent
from prompts import get_better_task_statement

agent = BasicAgent()


def get_answer_payload_results_log(questions_data: dict) -> tuple:
    """Get Answer Payload, Results Log or Error Message"""
    results_log = []
    answers_payload = []
    print(f"Running agent on {len(questions_data)} questions...")

    for item in questions_data:
        task_id = item.get("task_id")
        question_text = item.get("question")
        file_name = item.get("file_name")
        if not task_id or question_text is None:
            print(f"Skipping item with missing task_id or question: {item}")
            continue
        try:
            if file_name:
                task_statement = get_better_task_statement(question_text, file_name)
            else:
                task_statement = get_better_task_statement(question_text)
            submitted_answer = agent(task_statement)
            answers_payload.append(
                {"task_id": task_id, "submitted_answer": submitted_answer}
            )
            results_log.append(
                {
                    "Task ID": task_id,
                    "Question": question_text,
                    "Submitted Answer": submitted_answer,
                }
            )
        except Exception as e:  # pylint: disable=broad-exception-caught
            print(f"Error running agent on task {task_id}: {e}")
            results_log.append(
                {
                    "Task ID": task_id,
                    "Question": question_text,
                    "Submitted Answer": f"AGENT ERROR: {e}",
                }
            )
    results_df = pd.DataFrame(results_log)

    if not answers_payload:
        error_message = "Agent did not produce any answers to submit."
        print(error_message)
        return None, results_df, error_message

    return answers_payload, results_df, None