src/workflows/qb_agents.py

import time
from typing import Any, Iterable, TypedDict

from loguru import logger

from .executors import WorkflowOutput, execute_workflow
from .structs import TossupWorkflow, Workflow


def _get_workflow_response(
    workflow: Workflow, available_vars: dict[str, Any], logprob_step: bool | str = False
) -> tuple[WorkflowOutput, float]:
    """Get response from executing a complete workflow."""
    start_time = time.time()
    workflow_output = execute_workflow(workflow, available_vars, return_full_content=True, logprob_step=logprob_step)
    response_time = time.time() - start_time
    return workflow_output, response_time


class TossupResult(TypedDict):
    answer: str
    confidence: float
    logprob: float | None
    buzz: bool
    question_fragment: str
    position: int
    step_contents: list[str]
    response_time: float
    step_outputs: dict[str, Any]


class BonusResult(TypedDict):
    answer: str
    confidence: float
    explanation: str
    response_time: float
    step_contents: list[str]
    step_outputs: dict[str, Any]


class QuizBowlTossupAgent:
    """Agent for handling tossup questions with multiple steps in the workflow."""

    external_input_variable = "question_text"
    output_variables = ["answer", "confidence"]

    def __init__(self, workflow: TossupWorkflow):
        """Initialize the multi-step tossup agent.

        Args:
            workflow: The workflow containing multiple steps
            buzz_threshold: Confidence threshold for buzzing
        """
        self.workflow = workflow
        self.output_variables = list(workflow.outputs.keys())

        # Validate input variables
        if self.external_input_variable not in workflow.inputs:
            raise ValueError(f"External input variable {self.external_input_variable} not found in workflow inputs")

        # Validate output variables
        for out_var in self.output_variables:
            if out_var not in workflow.outputs:
                raise ValueError(f"Output variable {out_var} not found in workflow outputs")

    def _single_run(self, question_run: str, position: int) -> TossupResult:
        """Process a single question run."""
        answer_var_step = self.workflow.outputs["answer"].split(".")[0]
        workflow_output, response_time = _get_workflow_response(
            self.workflow, {self.external_input_variable: question_run}, logprob_step=answer_var_step
        )
        final_outputs = workflow_output["final_outputs"]
        buzz = self.workflow.buzzer.run(final_outputs["confidence"], logprob=workflow_output["logprob"])
        result: TossupResult = {
            "position": position,
            "answer": final_outputs["answer"],
            "confidence": final_outputs["confidence"],
            "logprob": workflow_output["logprob"],
            "buzz": buzz,
            "question_fragment": question_run,
            "step_contents": workflow_output["step_contents"],
            "step_outputs": workflow_output["intermediate_outputs"],  # Include intermediate step outputs
            "response_time": response_time,
        }
        return result

    def run(self, question_runs: list[str], early_stop: bool = True) -> Iterable[TossupResult]:
        """Process a tossup question and decide when to buzz based on confidence.

        Args:
            question_runs: Progressive reveals of the question text
            early_stop: Whether to stop after the first buzz

        Yields:
            Dict containing:
                - answer: The model's answer
                - confidence: Confidence score
                - buzz: Whether to buzz
                - question_fragment: Current question text
                - position: Current position in question
                - step_contents: String content outputs of each step
                - response_time: Time taken for response
                - step_outputs: Outputs from each step
        """
        for i, question_text in enumerate(question_runs):
            # Execute the complete workflow
            result = self._single_run(question_text, i + 1)

            yield result

            # If we've reached the confidence threshold, buzz and stop
            if early_stop and result["buzz"]:
                if i + 1 < len(question_runs):
                    yield self._single_run(question_runs[-1], len(question_runs))
                return


class QuizBowlBonusAgent:
    """Agent for handling bonus questions with multiple steps in the workflow."""

    external_input_variables = ["leadin", "part"]
    output_variables = ["answer", "confidence", "explanation"]

    def __init__(self, workflow: Workflow):
        """Initialize the multi-step bonus agent.

        Args:
            workflow: The workflow containing multiple steps
        """
        self.workflow = workflow
        self.output_variables = list(workflow.outputs.keys())

        # Validate input variables
        for input_var in self.external_input_variables:
            if input_var not in workflow.inputs:
                raise ValueError(f"External input variable {input_var} not found in workflow inputs")

        # Validate output variables
        for out_var in self.output_variables:
            if out_var not in workflow.outputs:
                raise ValueError(f"Output variable {out_var} not found in workflow outputs")

    def run(self, leadin: str, part: str) -> BonusResult:
        """Process a bonus part with the given leadin.

        Args:
            leadin: The leadin text for the bonus question
            part: The specific part text to answer

        Returns:
            Dict containing:
                - answer: The model's answer
                - confidence: Confidence score
                - explanation: Explanation for the answer
                - step_contents: String content outputs of each step
                - response_time: Time taken for response
                - step_outputs: Outputs from each step
        """
        workflow_output, response_time = _get_workflow_response(
            self.workflow,
            {
                "leadin": leadin,
                "part": part,
            },
        )
        final_outputs = workflow_output["final_outputs"]
        return {
            "answer": final_outputs["answer"],
            "confidence": final_outputs["confidence"],
            "explanation": final_outputs["explanation"],
            "step_contents": workflow_output["step_contents"],
            "response_time": response_time,
            "step_outputs": workflow_output["intermediate_outputs"],  # Include intermediate step outputs
        }


# Example usage
if __name__ == "__main__":
    # Load the Quizbowl dataset
    from datasets import load_dataset

    from workflows.factory import create_quizbowl_bonus_workflow, create_quizbowl_tossup_workflow

    ds_name = "umdclip/leaderboard_co_set"
    ds = load_dataset(ds_name, split="train")

    # Create the agents with multi-step workflows
    tossup_workflow = create_quizbowl_tossup_workflow()
    tossup_agent = QuizBowlTossupAgent(workflow=tossup_workflow, buzz_threshold=0.9)

    bonus_workflow = create_quizbowl_bonus_workflow()
    bonus_agent = QuizBowlBonusAgent(workflow=bonus_workflow)

    # Example for tossup mode
    print("\n=== TOSSUP MODE EXAMPLE ===")
    sample_question = ds[30]
    print(sample_question["question_runs"][-1])
    print(sample_question["gold_label"])
    print()
    question_runs = sample_question["question_runs"]

    results = tossup_agent.run(question_runs, early_stop=True)
    for result in results:
        print(result["step_contents"])
        print(f"Guess at position {result['position']}: {result['answer']}")
        print(f"Confidence: {result['confidence']}")
        print("Step outputs:", result["step_outputs"])
        if result["buzz"]:
            print("Buzzed!\n")

    # Example for bonus mode
    print("\n=== BONUS MODE EXAMPLE ===")
    sample_bonus = ds[31]  # Assuming this is a bonus question
    leadin = sample_bonus["leadin"]
    parts = sample_bonus["parts"]

    print(f"Leadin: {leadin}")
    for i, part in enumerate(parts):
        print(f"\nPart {i + 1}: {part['part']}")
        result = bonus_agent.run(leadin, part["part"])
        print(f"Answer: {result['answer']}")
        print(f"Confidence: {result['confidence']}")
        print(f"Explanation: {result['explanation']}")
        print(f"Response time: {result['response_time']:.2f}s")
        print("Step outputs:", result["step_outputs"])