Session

def add_round(self, id: int, round: BaseRound) -> None:
    """
    Add a round to the session.
    :param id: The id of the round.
    :param round: The round to be added.
    """
    self._rounds[id] = round

async def capture_last_screenshot(
    self, save_path: str, full_screen: bool = False
) -> None:
    """
    Capture the last window screenshot.
    :param save_path: The path to save the window screenshot.
    :param full_screen: Whether to capture the full screen or just the active window.
    """

    try:
        if full_screen:
            command = Command(
                tool_name="capture_desktop_screenshot",
                parameters={"all_screens": True},
                tool_type="data_collection",
            )
        else:

            command = Command(
                tool_name="capture_window_screenshot",
                parameters={},
                tool_type="data_collection",
            )

        result = await self.context.command_dispatcher.execute_commands([command])
        image = result[0].result

        self.logger.info(f"Captured screenshot at final: {save_path}")
        if image:
            utils.save_image_string(image, save_path)

    except Exception as e:
        self.logger.warning(
            f"The last snapshot capture failed, due to the error: {e}"
        )

async def capture_last_snapshot(self) -> None:
    """
    Capture the last snapshot of the application, including the screenshot and the XML file if configured.
    """  # Capture the final screenshot
    screenshot_save_path = self.log_path + "action_step_final.png"

    if (
        self.application_window is not None
        or self.application_window_info is not None
    ):

        await self.capture_last_screenshot(screenshot_save_path)

        if ufo_config.system.save_ui_tree:
            ui_tree_path = os.path.join(self.log_path, "ui_trees")
            ui_tree_file_name = "ui_tree_final.json"
            ui_tree_save_path = os.path.join(ui_tree_path, ui_tree_file_name)
            await self.capture_last_ui_tree(ui_tree_save_path)

        if ufo_config.system.save_full_screen:

            desktop_save_path = self.log_path + "desktop_final.png"

            await self.capture_last_screenshot(desktop_save_path, full_screen=True)

async def capture_last_ui_tree(self, save_path: str) -> None:
    """
    Capture the last UI tree snapshot.
    :param save_path: The path to save the UI tree snapshot.
    """

    result = await self.context.command_dispatcher.execute_commands(
        [
            Command(
                tool_name="get_ui_tree",
                parameters={},
                tool_type="data_collection",
            )
        ]
    )

    if result and result[0].result:
        with open(save_path, "w") as file:
            json.dump(result[0].result, file, indent=4)

def create_following_round(self) -> BaseRound:
    """
    Create a following round.
    return: The following round.
    """
    pass

@abstractmethod
def create_new_round(self) -> Optional[BaseRound]:
    """
    Create a new round.
    """
    pass

def evaluation(self) -> None:
    """
    Evaluate the session.
    """
    console.print("📊 Evaluating the session...", style="yellow")

    is_visual = ufo_config.evaluation_agent.visual_mode

    evaluator = EvaluationAgent(
        name="eva_agent",
        is_visual=is_visual,
        main_prompt=ufo_config.system.EVALUATION_PROMPT,
        example_prompt="",
    )

    requests = self.request_to_evaluate()

    # Evaluate the session, first use the default setting, if failed, then disable the screenshot evaluation.
    try:
        result, cost = evaluator.evaluate(
            request=requests,
            log_path=self.log_path,
            eva_all_screenshots=ufo_config.system.eva_all_screenshots,
            context=self.context,
        )
    except Exception as e:
        result, cost = evaluator.evaluate(
            request=requests,
            log_path=self.log_path,
            eva_all_screenshots=False,
            context=self.context,
        )

    # Add additional information to the evaluation result.
    additional_info = {
        "level": "session",
        "request": requests,
        "type": "evaluation_result",
    }
    result.update(additional_info)

    self._results.append(result)

    self.cost += cost

    evaluator.print_response(result)

    self.evaluation_logger.write(json.dumps(result))

    self.logger.info(
        f"Evaluation result saved to {os.path.join(self.log_path, 'evaluation.log')}"
    )

def experience_saver(self) -> None:
    """
    Save the current trajectory as agent experience.
    """
    console.print(
        "📚 Summarizing and saving the execution flow as experience...",
        style="yellow",
    )

    summarizer = ExperienceSummarizer(
        ufo_config.app_agent.visual_mode,
        ufo_config.system.EXPERIENCE_PROMPT,
        ufo_config.system.APPAGENT_EXAMPLE_PROMPT,
        ufo_config.system.API_PROMPT,
    )
    experience = summarizer.read_logs(self.log_path)
    summaries, cost = summarizer.get_summary_list(experience)

    experience_path = ufo_config.system.EXPERIENCE_SAVED_PATH
    utils.create_folder(experience_path)
    summarizer.create_or_update_yaml(
        summaries, os.path.join(experience_path, "experience.yaml")
    )
    summarizer.create_or_update_vector_db(
        summaries, os.path.join(experience_path, "experience_db")
    )

    self.cost += cost
    self.logger.info(f"The experience has been saved to {experience_path}")

@staticmethod
def initialize_logger(log_path: str, log_filename: str, mode="a") -> logging.Logger:
    """
    Initialize logging.
    log_path: The path of the log file.
    log_filename: The name of the log file.
    return: The logger.
    """
    # Code for initializing logging
    logger = logging.Logger(log_filename)

    if not ufo_config.system.print_log:
        # Remove existing handlers if PRINT_LOG is False
        logger.handlers = []

    log_file_path = os.path.join(log_path, log_filename)
    file_handler = logging.FileHandler(log_file_path, mode=mode, encoding="utf-8")
    formatter = logging.Formatter("%(message)s")
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    logger.setLevel(ufo_config.system.log_level)

    return logger

def is_error(self):
    """
    Check if the session is in error state.
    return: True if the session is in error state, otherwise False.
    """
    if self.current_round is not None:
        return self.current_round.state.name() == AgentStatus.ERROR.value
    return False

def is_finished(self) -> bool:
    """
    Check if the session is ended.
    return: True if the session is ended, otherwise False.
    """
    if (
        self._finish
        or self.step >= ufo_config.system.max_step
        or self.total_rounds >= ufo_config.system.max_round
    ):
        return True

    if self.is_error():
        return True

    return False

@abstractmethod
def next_request(self) -> str:
    """
    Get the next request of the session.
    return: The request of the session.
    """
    pass

def print_cost(self) -> None:
    """
    Print the total cost of the session.
    """

    if isinstance(self.cost, float) and self.cost > 0:
        formatted_cost = "${:.2f}".format(self.cost)
        console.print(
            f"💰 Total request cost of the session: {formatted_cost}",
            style="yellow",
        )
    else:
        console.print(
            f"ℹ️  Cost is not available for the model {ufo_config.host_agent.api_model} or {ufo_config.app_agent.api_model}.",
            style="yellow",
        )
        self.logger.warning("Cost information is not available.")

@abstractmethod
def request_to_evaluate(self) -> str:
    """
    Get the request to evaluate.
    return: The request(s) to evaluate.
    """
    pass

@abstractmethod
def reset(self) -> None:
    """
    Reset the session to initial state.
    """
    pass

async def run(self) -> List[Dict[str, str]]:
    """
    Run the session.
    :return: The result per session
    """

    while not self.is_finished():

        round = self.create_new_round()
        if round is None:
            break

        round_result = await round.run()

        self.results.append({"request": round.request, "result": round_result})

    await self.capture_last_snapshot()

    if self._should_evaluate and not self.is_error():
        self.evaluation()

    if ufo_config.system.log_to_markdown:

        self.save_log_to_markdown()

    self.print_cost()

    return self.results

def save_log_to_markdown(self) -> None:
    """
    Save the log of the session to markdown file.
    """

    file_path = self.log_path
    trajectory = Trajectory(file_path)
    trajectory.to_markdown(file_path + "/output.md")
    self.logger.info(f"Trajectory saved to {file_path + '/output.md'}")

def create_new_round(self) -> Optional[BaseRound]:
    """
    Create a new round.
    """

    # Get a request for the new round.
    request = self.next_request()

    # Create a new round and return None if the session is finished.

    if self.is_finished():
        return None

    self._host_agent.set_state(self._host_agent.default_state)

    round = BaseRound(
        request=request,
        agent=self._host_agent,
        context=self.context,
        should_evaluate=ufo_config.system.eva_round,
        id=self.total_rounds,
    )

    self.add_round(round.id, round)

    return round

def next_request(self) -> str:
    """
    Get the request for the host agent.
    :return: The request for the host agent.
    """
    if self.total_rounds == 0:

        # If the request is provided via command line, use it directly.
        if self._init_request:
            return self._init_request
        # Otherwise, ask the user to input the request with enhanced UX.
        else:
            return interactor.first_request()
    else:
        request, iscomplete = interactor.new_request()
        if iscomplete:
            self._finish = True
        return request

def request_to_evaluate(self) -> str:
    """
    Get the request to evaluate.
    return: The request(s) to evaluate.
    """
    request_memory = self._host_agent.blackboard.requests
    return request_memory.to_json()

async def run(self) -> None:
    """
    Run the session.
    """
    await super().run()
    # Save the experience if the user asks so.

    save_experience = ufo_config.system.save_experience

    self.logger.info(f"Save experience setting: {save_experience}")

    if save_experience == "always":
        self.experience_saver()
    elif save_experience == "ask":
        if interactor.experience_asker():
            self.experience_saver()

    elif save_experience == "auto":
        task_completed = self.results.get("complete", "no")
        if task_completed.lower() == "yes":
            self.experience_saver()

    elif save_experience == "always_not":
        pass

def create_new_round(self) -> Optional[BaseRound]:
    """
    Create a new round for Linux session.
    Since there's no host agent, directly create app-level rounds.
    """
    request = self.next_request()

    if self.is_finished():
        return None

    round = BaseRound(
        request=request,
        agent=self._agent,
        context=self.context,
        should_evaluate=ufo_config.system.eva_round,
        id=self.total_rounds,
    )

    self.add_round(round.id, round)
    return round

def next_request(self) -> str:
    """
    Get the request for the app agent.
    :return: The request for the app agent.
    """
    if self.total_rounds == 0:
        if self._init_request:
            return self._init_request
        else:
            return interactor.first_request()
    else:
        request, iscomplete = interactor.new_request()
        if iscomplete:
            self._finish = True
        return request

def request_to_evaluate(self) -> str:
    """
    Get the request to evaluate.
    :return: The request(s) to evaluate.
    """
    # For Linux session, collect requests from all rounds
    if self.current_round and hasattr(self.current_round.agent, "blackboard"):
        request_memory = self.current_round.agent.blackboard.requests
        return request_memory.to_json()
    return self._init_request