{code(think)}"""
    return action_info

def get_state_error(state: Info):
    try:
        step_info = state.exp_result.steps_info[state.step + 1]
        err_msg = step_info.obs.get("last_action_error", None)
    except (IndexError, AttributeError):
        err_msg = None

    if err_msg is None or len(err_msg) == 0:
        err_msg = "No Error"
    return f"""\
**Step error after action:**

{code(err_msg)}"""

def get_seeds_df(result_df: pd.DataFrame, task_name: str):
    result_df = result_df.reset_index(inplace=False)
    result_df = result_df[result_df[TASK_NAME_KEY] == task_name]

    def extract_columns(row: pd.Series):
        return pd.Series(
            {
                "idx": row.get("_row_index", None),
                "seed": row.get(TASK_SEED_KEY, None),
                "reward": row.get("cum_reward", None),
                "err": bool(row.get("err_msg", None)),
                "n_steps": row.get("n_steps", None),
            }
        )

    seed_df = result_df.apply(extract_columns, axis=1)
    # Ensure column order and readability
    seed_df = seed_df[["seed", "reward", "err", "n_steps", "idx"]]
    return seed_df

def on_select_agent(evt: gr.SelectData, df: pd.DataFrame):
    # TODO try to find a clever way to solve the sort bug here
    return info.get_agent_id(df.iloc[evt.index[0]])

def on_select_task(evt: gr.SelectData, df: pd.DataFrame, agent_id: list[tuple]):
    # get col index
    col_idx = df.columns.get_loc(TASK_NAME_KEY)
    return (agent_id, evt.row_value[col_idx])

def update_seeds(agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    seed_df = get_seeds_df(info.agent_df, task_name)
    first_seed = int(seed_df.iloc[0]["seed"])
    first_index = int(seed_df.iloc[0]["idx"])
    return seed_df, EpisodeId(
        agent_id=agent_id, task_name=task_name, seed=first_seed, row_index=first_index


    result_df = result_df.reset_index(inplace=False)
    result_df = result_df[result_df[TASK_NAME_KEY] == task_name]

    def extract_columns(row: pd.Series):
        return pd.Series(
            {
                "idx": row.get("_row_index", None),
                "seed": row.get(TASK_SEED_KEY, None),
                "reward": row.get("cum_reward", None),
                "err": bool(row.get("err_msg", None)),
                "n_steps": row.get("n_steps", None),
            }
        )

    seed_df = result_df.apply(extract_columns, axis=1)
    # Ensure column order and readability
    seed_df = seed_df[["seed", "reward", "err", "n_steps", "idx"]]
    return seed_df

def on_select_agent(evt: gr.SelectData, df: pd.DataFrame):
    # TODO try to find a clever way to solve the sort bug here
    return info.get_agent_id(df.iloc[evt.index[0]])

def on_select_task(evt: gr.SelectData, df: pd.DataFrame, agent_id: list[tuple]):
    # get col index
    col_idx = df.columns.get_loc(TASK_NAME_KEY)
    return (agent_id, evt.row_value[col_idx])

def update_seeds(agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    seed_df = get_seeds_df(info.agent_df, task_name)
    first_seed = int(seed_df.iloc[0]["seed"])
    first_index = int(seed_df.iloc[0]["idx"])
    return seed_df, EpisodeId(
        agent_id=agent_id, task_name=task_name, seed=first_seed, row_index=first_index
    )

def on_select_seed(evt: gr.SelectData, df: pd.DataFrame, agent_task_id: tuple):
    agent_id, task_name = agent_task_id


            {
                "idx": row.get("_row_index", None),
                "seed": row.get(TASK_SEED_KEY, None),
                "reward": row.get("cum_reward", None),
                "err": bool(row.get("err_msg", None)),
                "n_steps": row.get("n_steps", None),
            }
        )

    seed_df = result_df.apply(extract_columns, axis=1)
    # Ensure column order and readability
    seed_df = seed_df[["seed", "reward", "err", "n_steps", "idx"]]
    return seed_df

def on_select_agent(evt: gr.SelectData, df: pd.DataFrame):
    # TODO try to find a clever way to solve the sort bug here
    return info.get_agent_id(df.iloc[evt.index[0]])

def on_select_task(evt: gr.SelectData, df: pd.DataFrame, agent_id: list[tuple]):
    # get col index
    col_idx = df.columns.get_loc(TASK_NAME_KEY)
    return (agent_id, evt.row_value[col_idx])

def update_seeds(agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    seed_df = get_seeds_df(info.agent_df, task_name)
    first_seed = int(seed_df.iloc[0]["seed"])
    first_index = int(seed_df.iloc[0]["idx"])
    return seed_df, EpisodeId(
        agent_id=agent_id, task_name=task_name, seed=first_seed, row_index=first_index
    )

def on_select_seed(evt: gr.SelectData, df: pd.DataFrame, agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    col_idx = df.columns.get_loc("seed")
    idx_col = df.columns.get_loc("idx")
    seed = evt.row_value[col_idx]
    row_index = evt.row_value[idx_col]
    return EpisodeId(agent_id=agent_id, task_name=task_name, seed=seed, row_index=row_index)


            }
        )

    seed_df = result_df.apply(extract_columns, axis=1)
    # Ensure column order and readability
    seed_df = seed_df[["seed", "reward", "err", "n_steps", "idx"]]
    return seed_df

def on_select_agent(evt: gr.SelectData, df: pd.DataFrame):
    # TODO try to find a clever way to solve the sort bug here
    return info.get_agent_id(df.iloc[evt.index[0]])

def on_select_task(evt: gr.SelectData, df: pd.DataFrame, agent_id: list[tuple]):
    # get col index
    col_idx = df.columns.get_loc(TASK_NAME_KEY)
    return (agent_id, evt.row_value[col_idx])

def update_seeds(agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    seed_df = get_seeds_df(info.agent_df, task_name)
    first_seed = int(seed_df.iloc[0]["seed"])
    first_index = int(seed_df.iloc[0]["idx"])
    return seed_df, EpisodeId(
        agent_id=agent_id, task_name=task_name, seed=first_seed, row_index=first_index
    )

def on_select_seed(evt: gr.SelectData, df: pd.DataFrame, agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    col_idx = df.columns.get_loc("seed")
    idx_col = df.columns.get_loc("idx")
    seed = evt.row_value[col_idx]
    row_index = evt.row_value[idx_col]
    return EpisodeId(agent_id=agent_id, task_name=task_name, seed=seed, row_index=row_index)

def new_episode(episode_id: EpisodeId, progress=gr.Progress()):
    print("new_episode", episode_id)
    global info
    info.update_exp_result(episode_id=episode_id)
    return generate_profiling(progress.tqdm), StepId(episode_id, info.step)

def fig_to_pil(fig):
    buf = BytesIO()


    # TODO try to find a clever way to solve the sort bug here
    return info.get_agent_id(df.iloc[evt.index[0]])

def on_select_task(evt: gr.SelectData, df: pd.DataFrame, agent_id: list[tuple]):
    # get col index
    col_idx = df.columns.get_loc(TASK_NAME_KEY)
    return (agent_id, evt.row_value[col_idx])

def update_seeds(agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    seed_df = get_seeds_df(info.agent_df, task_name)
    first_seed = int(seed_df.iloc[0]["seed"])
    first_index = int(seed_df.iloc[0]["idx"])
    return seed_df, EpisodeId(
        agent_id=agent_id, task_name=task_name, seed=first_seed, row_index=first_index
    )

def on_select_seed(evt: gr.SelectData, df: pd.DataFrame, agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    col_idx = df.columns.get_loc("seed")
    idx_col = df.columns.get_loc("idx")
    seed = evt.row_value[col_idx]
    row_index = evt.row_value[idx_col]
    return EpisodeId(agent_id=agent_id, task_name=task_name, seed=seed, row_index=row_index)

def new_episode(episode_id: EpisodeId, progress=gr.Progress()):
    print("new_episode", episode_id)
    global info
    info.update_exp_result(episode_id=episode_id)
    return generate_profiling(progress.tqdm), StepId(episode_id, info.step)

def fig_to_pil(fig):
    buf = BytesIO()
    fig.savefig(buf, format="png")
    buf.seek(0)
    img_pil = Image.open(buf)
    plt.close(fig)
    return img_pil

def format_constant_and_variables():
    global info


def update_seeds(agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    seed_df = get_seeds_df(info.agent_df, task_name)
    first_seed = int(seed_df.iloc[0]["seed"])
    first_index = int(seed_df.iloc[0]["idx"])
    return seed_df, EpisodeId(
        agent_id=agent_id, task_name=task_name, seed=first_seed, row_index=first_index
    )

def on_select_seed(evt: gr.SelectData, df: pd.DataFrame, agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    col_idx = df.columns.get_loc("seed")
    idx_col = df.columns.get_loc("idx")
    seed = evt.row_value[col_idx]
    row_index = evt.row_value[idx_col]
    return EpisodeId(agent_id=agent_id, task_name=task_name, seed=seed, row_index=row_index)

def new_episode(episode_id: EpisodeId, progress=gr.Progress()):
    print("new_episode", episode_id)
    global info
    info.update_exp_result(episode_id=episode_id)
    return generate_profiling(progress.tqdm), StepId(episode_id, info.step)

def fig_to_pil(fig):
    buf = BytesIO()
    fig.savefig(buf, format="png")
    buf.seek(0)
    img_pil = Image.open(buf)
    plt.close(fig)
    return img_pil

def format_constant_and_variables():
    global info
    df = info.result_df
    constants, variables, _ = inspect_results.get_constants_and_variables(df)

    # map constants, a dict to a 2 column data frame with name and value
    constants = pd.DataFrame(constants.items(), columns=["name", "value"])
    records = []
    for var in variables:


        agent_id=agent_id, task_name=task_name, seed=first_seed, row_index=first_index
    )

def on_select_seed(evt: gr.SelectData, df: pd.DataFrame, agent_task_id: tuple):
    agent_id, task_name = agent_task_id
    col_idx = df.columns.get_loc("seed")
    idx_col = df.columns.get_loc("idx")
    seed = evt.row_value[col_idx]
    row_index = evt.row_value[idx_col]
    return EpisodeId(agent_id=agent_id, task_name=task_name, seed=seed, row_index=row_index)

def new_episode(episode_id: EpisodeId, progress=gr.Progress()):
    print("new_episode", episode_id)
    global info
    info.update_exp_result(episode_id=episode_id)
    return generate_profiling(progress.tqdm), StepId(episode_id, info.step)

def fig_to_pil(fig):
    buf = BytesIO()
    fig.savefig(buf, format="png")
    buf.seek(0)
    img_pil = Image.open(buf)
    plt.close(fig)
    return img_pil

def format_constant_and_variables():
    global info
    df = info.result_df
    constants, variables, _ = inspect_results.get_constants_and_variables(df)

    # map constants, a dict to a 2 column data frame with name and value
    constants = pd.DataFrame(constants.items(), columns=["name", "value"])
    records = []
    for var in variables:
        if var == "stack_trace":
            continue

        # get unique with count and sort by count descending
        unique_counts = df[var].value_counts().sort_values(ascending=False)

        for i, (val, count) in enumerate(unique_counts.items()):
            record = {
                "Name": var,


    row_index = evt.row_value[idx_col]
    return EpisodeId(agent_id=agent_id, task_name=task_name, seed=seed, row_index=row_index)

def new_episode(episode_id: EpisodeId, progress=gr.Progress()):
    print("new_episode", episode_id)
    global info
    info.update_exp_result(episode_id=episode_id)
    return generate_profiling(progress.tqdm), StepId(episode_id, info.step)

def fig_to_pil(fig):
    buf = BytesIO()
    fig.savefig(buf, format="png")
    buf.seek(0)
    img_pil = Image.open(buf)
    plt.close(fig)
    return img_pil

def format_constant_and_variables():
    global info
    df = info.result_df
    constants, variables, _ = inspect_results.get_constants_and_variables(df)

    # map constants, a dict to a 2 column data frame with name and value
    constants = pd.DataFrame(constants.items(), columns=["name", "value"])
    records = []
    for var in variables:
        if var == "stack_trace":
            continue

        # get unique with count and sort by count descending
        unique_counts = df[var].value_counts().sort_values(ascending=False)

        for i, (val, count) in enumerate(unique_counts.items()):
            record = {
                "Name": var,
                "n unique": len(unique_counts),
                "i": i,
                "count": f"{count}/{len(df)}",
                "value": val,
            }

            records.append(record)
            if i >= 2:
                break

        if len(unique_counts) > 3:
            records.append(
                {
                    "Name": var,
                    "n unique": len(unique_counts),
                    "i": "...",
                    "count": "...",
                    "value": "...",
                }
            )
        records.append({"Name": ""})
    return constants, pd.DataFrame(records)

def get_agent_report(result_df: pd.DataFrame):
    levels = list(range(result_df.index.nlevels))

    if len(levels) == 1:
        result_df = result_df.set_index(AGENT_NAME_KEY, append=True)
        levels = list(range(result_df.index.nlevels))

    report = result_df.groupby(level=levels[1:]).apply(inspect_results.summarize)

    return report

def update_global_stats():
    try:
        stats = inspect_results.global_report(
            info.result_df, reduce_fn=inspect_results.summarize_stats
        )
        stats.reset_index(inplace=True)


            }

            records.append(record)
            if i >= 2:
                break

        if len(unique_counts) > 3:
            records.append(
                {
                    "Name": var,
                    "n unique": len(unique_counts),
                    "i": "...",
                    "count": "...",
                    "value": "...",
                }
            )
        records.append({"Name": ""})
    return constants, pd.DataFrame(records)

def get_agent_report(result_df: pd.DataFrame):
    levels = list(range(result_df.index.nlevels))

    if len(levels) == 1:
        result_df = result_df.set_index(AGENT_NAME_KEY, append=True)
        levels = list(range(result_df.index.nlevels))

    report = result_df.groupby(level=levels[1:]).apply(inspect_results.summarize)

    return report

def update_global_stats():
    try:
        stats = inspect_results.global_report(
            info.result_df, reduce_fn=inspect_results.summarize_stats
        )
        stats.reset_index(inplace=True)
        return stats

    except Exception as e:
        warning(f"Error while updating global stats: {e}")
        return None

def update_error_report():
    return inspect_results.error_report(info.result_df, max_stack_trace=3, use_log=True)

def new_exp_dir(study_names: list, progress=gr.Progress(), just_refresh=False):


                    "count": "...",
                    "value": "...",
                }
            )
        records.append({"Name": ""})
    return constants, pd.DataFrame(records)

def get_agent_report(result_df: pd.DataFrame):
    levels = list(range(result_df.index.nlevels))

    if len(levels) == 1:
        result_df = result_df.set_index(AGENT_NAME_KEY, append=True)
        levels = list(range(result_df.index.nlevels))

    report = result_df.groupby(level=levels[1:]).apply(inspect_results.summarize)

    return report

def update_global_stats():
    try:
        stats = inspect_results.global_report(
            info.result_df, reduce_fn=inspect_results.summarize_stats
        )
        stats.reset_index(inplace=True)
        return stats

    except Exception as e:
        warning(f"Error while updating global stats: {e}")
        return None

def update_error_report():
    return inspect_results.error_report(info.result_df, max_stack_trace=3, use_log=True)

def new_exp_dir(study_names: list, progress=gr.Progress(), just_refresh=False):
    global info

    # remove select_dir_instructions from study_names
    if select_dir_instructions in study_names:
        study_names.remove(select_dir_instructions)

    if len(study_names) == 0:
        return None, None, None, None, None, None

    info.study_dirs = [info.results_dir / study_name.split(" - ")[0] for study_name in study_names]
    info.result_df = inspect_results.load_result_df(info.study_dirs, progress_fn=progress.tqdm)
    info.result_df = remove_args_from_col(info.result_df)


        levels = list(range(result_df.index.nlevels))

    report = result_df.groupby(level=levels[1:]).apply(inspect_results.summarize)

    return report

def update_global_stats():
    try:
        stats = inspect_results.global_report(
            info.result_df, reduce_fn=inspect_results.summarize_stats
        )
        stats.reset_index(inplace=True)
        return stats

    except Exception as e:
        warning(f"Error while updating global stats: {e}")
        return None

def update_error_report():
    return inspect_results.error_report(info.result_df, max_stack_trace=3, use_log=True)

def new_exp_dir(study_names: list, progress=gr.Progress(), just_refresh=False):
    global info

    # remove select_dir_instructions from study_names
    if select_dir_instructions in study_names:
        study_names.remove(select_dir_instructions)

    if len(study_names) == 0:
        return None, None, None, None, None, None

    info.study_dirs = [info.results_dir / study_name.split(" - ")[0] for study_name in study_names]
    info.result_df = inspect_results.load_result_df(info.study_dirs, progress_fn=progress.tqdm)
    info.result_df = remove_args_from_col(info.result_df)

    study_summary = inspect_results.summarize_study(info.result_df)
    # save study_summary

    for study_dir in info.study_dirs:


    return report

def update_global_stats():
    try:
        stats = inspect_results.global_report(
            info.result_df, reduce_fn=inspect_results.summarize_stats
        )
        stats.reset_index(inplace=True)
        return stats

    except Exception as e:
        warning(f"Error while updating global stats: {e}")
        return None

def update_error_report():
    return inspect_results.error_report(info.result_df, max_stack_trace=3, use_log=True)

def new_exp_dir(study_names: list, progress=gr.Progress(), just_refresh=False):
    global info

    # remove select_dir_instructions from study_names
    if select_dir_instructions in study_names:
        study_names.remove(select_dir_instructions)

    if len(study_names) == 0:
        return None, None, None, None, None, None

    info.study_dirs = [info.results_dir / study_name.split(" - ")[0] for study_name in study_names]
    info.result_df = inspect_results.load_result_df(info.study_dirs, progress_fn=progress.tqdm)
    info.result_df = remove_args_from_col(info.result_df)

    study_summary = inspect_results.summarize_study(info.result_df)
    # save study_summary

    for study_dir in info.study_dirs:
        study_summary.to_csv(study_dir / "summary_df.csv", index=False)
        agent_report = display_table(study_summary)

    info.agent_id_keys = agent_report.index.names
    agent_report.reset_index(inplace=True)

    agent_id = info.get_agent_id(agent_report.iloc[0])

    constants, variables = format_constant_and_variables()
    return (
        agent_report,
        agent_id,
        constants,
        variables,
        update_global_stats(),
        update_error_report(),
    )

def new_agent_id(agent_id: list[tuple]):
    global info
    info.filter_agent_id(agent_id=agent_id)

    info.tasks_df = inspect_results.reduce_episodes(info.agent_df).reset_index()
    info.tasks_df = info.tasks_df.drop(columns=["std_err"])

    # task name of first element
    task_name = info.tasks_df.iloc[0][TASK_NAME_KEY]
    return info.tasks_df, (agent_id, task_name)

def get_directory_contents(results_dir: Path):
    exp_descriptions = []
    for dir in results_dir.iterdir():
        if not dir.is_dir():
            continue


    for study_dir in info.study_dirs:
        study_summary.to_csv(study_dir / "summary_df.csv", index=False)
        agent_report = display_table(study_summary)

    info.agent_id_keys = agent_report.index.names
    agent_report.reset_index(inplace=True)

    agent_id = info.get_agent_id(agent_report.iloc[0])

    constants, variables = format_constant_and_variables()
    return (
        agent_report,
        agent_id,
        constants,
        variables,
        update_global_stats(),
        update_error_report(),
    )

def new_agent_id(agent_id: list[tuple]):
    global info
    info.filter_agent_id(agent_id=agent_id)

    info.tasks_df = inspect_results.reduce_episodes(info.agent_df).reset_index()
    info.tasks_df = info.tasks_df.drop(columns=["std_err"])

    # task name of first element
    task_name = info.tasks_df.iloc[0][TASK_NAME_KEY]
    return info.tasks_df, (agent_id, task_name)

def get_directory_contents(results_dir: Path):
    exp_descriptions = []
    for dir in results_dir.iterdir():
        if not dir.is_dir():
            continue

        exp_description = dir.name
        try:
            # get summary*.csv files and find the most recent
            summary_files = list(dir.glob("summary*.csv"))
            if len(summary_files) != 0:
                most_recent_summary = max(summary_files, key=os.path.getctime)
                summary_df = pd.read_csv(most_recent_summary)

                if len(summary_df) == 0:
                    continue  # skip if all avg_reward are NaN

                # get row with max avg_reward


        agent_id,
        constants,
        variables,
        update_global_stats(),
        update_error_report(),
    )

def new_agent_id(agent_id: list[tuple]):
    global info
    info.filter_agent_id(agent_id=agent_id)

    info.tasks_df = inspect_results.reduce_episodes(info.agent_df).reset_index()
    info.tasks_df = info.tasks_df.drop(columns=["std_err"])

    # task name of first element
    task_name = info.tasks_df.iloc[0][TASK_NAME_KEY]
    return info.tasks_df, (agent_id, task_name)

def get_directory_contents(results_dir: Path):
    exp_descriptions = []
    for dir in results_dir.iterdir():
        if not dir.is_dir():
            continue

        exp_description = dir.name
        try:
            # get summary*.csv files and find the most recent
            summary_files = list(dir.glob("summary*.csv"))
            if len(summary_files) != 0:
                most_recent_summary = max(summary_files, key=os.path.getctime)
                summary_df = pd.read_csv(most_recent_summary)

                if len(summary_df) == 0:
                    continue  # skip if all avg_reward are NaN

                # get row with max avg_reward
                max_reward_row = summary_df.loc[summary_df["avg_reward"].idxmax(skipna=True)]
                reward = max_reward_row["avg_reward"] * 100
                completed = max_reward_row["n_completed"]
                n_err = max_reward_row["n_err"]
                exp_description += (
                    f" - avg-reward: {reward:.1f}% - completed: {completed} - errors: {n_err}"
                )
        except Exception as e:
            print(f"Error while reading summary file {most_recent_summary}: {e}")

        exp_descriptions.append(exp_description)

    return [select_dir_instructions] + sorted(exp_descriptions, reverse=True)

def most_recent_folder(results_dir: Path):
    return get_most_recent_study(results_dir).name

def refresh_exp_dir_choices(exp_dir_choice):
    global info
    return gr.Dropdown(
        choices=get_directory_contents(info.results_dir), value=exp_dir_choice, scale=1
    )

def generate_profiling(progress_fn):
    global info

    if info.exp_result is None:
        return None

    fig, ax = plt.subplots(figsize=(20, 3))


                if len(summary_df) == 0:
                    continue  # skip if all avg_reward are NaN

                # get row with max avg_reward
                max_reward_row = summary_df.loc[summary_df["avg_reward"].idxmax(skipna=True)]
                reward = max_reward_row["avg_reward"] * 100
                completed = max_reward_row["n_completed"]
                n_err = max_reward_row["n_err"]
                exp_description += (
                    f" - avg-reward: {reward:.1f}% - completed: {completed} - errors: {n_err}"
                )
        except Exception as e:
            print(f"Error while reading summary file {most_recent_summary}: {e}")

        exp_descriptions.append(exp_description)

    return [select_dir_instructions] + sorted(exp_descriptions, reverse=True)

def most_recent_folder(results_dir: Path):
    return get_most_recent_study(results_dir).name

def refresh_exp_dir_choices(exp_dir_choice):
    global info
    return gr.Dropdown(
        choices=get_directory_contents(info.results_dir), value=exp_dir_choice, scale=1
    )

def generate_profiling(progress_fn):
    global info

    if info.exp_result is None:
        return None

    fig, ax = plt.subplots(figsize=(20, 3))

    try:
        summary_info = info.exp_result.summary_info
    except FileNotFoundError:


                # get row with max avg_reward
                max_reward_row = summary_df.loc[summary_df["avg_reward"].idxmax(skipna=True)]
                reward = max_reward_row["avg_reward"] * 100
                completed = max_reward_row["n_completed"]
                n_err = max_reward_row["n_err"]
                exp_description += (
                    f" - avg-reward: {reward:.1f}% - completed: {completed} - errors: {n_err}"
                )
        except Exception as e:
            print(f"Error while reading summary file {most_recent_summary}: {e}")

        exp_descriptions.append(exp_description)

    return [select_dir_instructions] + sorted(exp_descriptions, reverse=True)

def most_recent_folder(results_dir: Path):
    return get_most_recent_study(results_dir).name

def refresh_exp_dir_choices(exp_dir_choice):
    global info
    return gr.Dropdown(
        choices=get_directory_contents(info.results_dir), value=exp_dir_choice, scale=1
    )

def generate_profiling(progress_fn):
    global info

    if info.exp_result is None:
        return None

    fig, ax = plt.subplots(figsize=(20, 3))

    try:
        summary_info = info.exp_result.summary_info
    except FileNotFoundError:
        summary_info = {}

    info.exp_result.progress_fn = progress_fn
    steps_info = info.exp_result.steps_info
    info.exp_result.progress_fn = None

    step_times = plot_profiling(ax, steps_info, summary_info, progress_fn)


                )
        except Exception as e:
            print(f"Error while reading summary file {most_recent_summary}: {e}")

        exp_descriptions.append(exp_description)

    return [select_dir_instructions] + sorted(exp_descriptions, reverse=True)

def most_recent_folder(results_dir: Path):
    return get_most_recent_study(results_dir).name

def refresh_exp_dir_choices(exp_dir_choice):
    global info
    return gr.Dropdown(
        choices=get_directory_contents(info.results_dir), value=exp_dir_choice, scale=1
    )

def generate_profiling(progress_fn):
    global info

    if info.exp_result is None:
        return None

    fig, ax = plt.subplots(figsize=(20, 3))

    try:
        summary_info = info.exp_result.summary_info
    except FileNotFoundError:
        summary_info = {}

    info.exp_result.progress_fn = progress_fn
    steps_info = info.exp_result.steps_info
    info.exp_result.progress_fn = None

    step_times = plot_profiling(ax, steps_info, summary_info, progress_fn)
    fig.tight_layout()
    info.click_mapper = ClickMapper(ax, step_times=step_times)

    return fig_to_pil(fig)

def add_patch(ax, start, stop, color, label, edge=False):
    if edge:
        ax.add_patch(
            patches.Rectangle(
                (start, 0),
                stop - start,
                1,
                edgecolor=color,
                alpha=1,
                label=label,
                fill=False,
                linewidth=3,
            )
        )
    else:
        ax.add_patch(
            patches.Rectangle((start, 0), stop - start, 1, color=color, alpha=1, label=label)
        )


        return None

    fig, ax = plt.subplots(figsize=(20, 3))

    try:
        summary_info = info.exp_result.summary_info
    except FileNotFoundError:
        summary_info = {}

    info.exp_result.progress_fn = progress_fn
    steps_info = info.exp_result.steps_info
    info.exp_result.progress_fn = None

    step_times = plot_profiling(ax, steps_info, summary_info, progress_fn)
    fig.tight_layout()
    info.click_mapper = ClickMapper(ax, step_times=step_times)

    return fig_to_pil(fig)

def add_patch(ax, start, stop, color, label, edge=False):
    if edge:
        ax.add_patch(
            patches.Rectangle(
                (start, 0),
                stop - start,
                1,
                edgecolor=color,
                alpha=1,
                label=label,
                fill=False,
                linewidth=3,
            )
        )
    else:
        ax.add_patch(
            patches.Rectangle((start, 0), stop - start, 1, color=color, alpha=1, label=label)
        )

def plot_profiling(ax, step_info_list: list[StepInfo], summary_info: dict, progress_fn):
    if len(step_info_list) == 0:
        warning("No step info to plot")
        return None

    # Updated labels to include new profiling stages
    labels = [
        "reset",
        "env",
        "agent",
        "exec action",
        "action error",
        "wait for page",
        "validation",
        "get observation",
    ]
    labels = {e: e for e in labels}


def add_patch(ax, start, stop, color, label, edge=False):
    if edge:
        ax.add_patch(
            patches.Rectangle(
                (start, 0),
                stop - start,
                1,
                edgecolor=color,
                alpha=1,
                label=label,
                fill=False,
                linewidth=3,
            )
        )
    else:
        ax.add_patch(
            patches.Rectangle((start, 0), stop - start, 1, color=color, alpha=1, label=label)
        )

def plot_profiling(ax, step_info_list: list[StepInfo], summary_info: dict, progress_fn):
    if len(step_info_list) == 0:
        warning("No step info to plot")
        return None

    # Updated labels to include new profiling stages
    labels = [
        "reset",
        "env",
        "agent",
        "exec action",
        "action error",
        "wait for page",
        "validation",
        "get observation",
    ]
    labels = {e: e for e in labels}

    colors = plt.get_cmap("tab20c").colors

    t0 = step_info_list[0].profiling.env_start
    all_times = []
    step_times = []
    for i, step_info in progress_fn(list(enumerate(step_info_list)), desc="Building plot."):
        assert isinstance(
            step_info, (StepInfo, BGymStepInfo)
        ), f"Expected StepInfo or BGymStepInfo, got {type(step_info)}"
        step = step_info.step

        prof = deepcopy(step_info.profiling)
        # remove t0 from elements in profiling using for
        for key, value in prof.__dict__.items():
            if isinstance(value, float):
                setattr(prof, key, value - t0)
                all_times.append(value - t0)

        if i == 0:
            # reset
            add_patch(ax, prof.env_start, prof.env_stop, colors[14], labels.pop("reset", None))

        else:
            # env
            add_patch(ax, prof.env_start, prof.env_stop, colors[1], labels.pop("env", None))

            # action
            label = labels.pop("exec action", None)
            add_patch(ax, prof.action_exec_start, prof.action_exec_stop, colors[3], label)

            # NEW: Add wait for page loading visualization
            if (
                hasattr(prof, "wait_for_page_loading_start")
                and prof.wait_for_page_loading_start is not None
                and prof.wait_for_page_loading_start > 0
            ):
                add_patch(
                    ax,
                    prof.wait_for_page_loading_start,
                    prof.wait_for_page_loading_stop,
                    colors[19],
                    labels.pop("wait for page", None),
                )

            # NEW: Add validation visualization
            if (
                hasattr(prof, "validation_start")
                and prof.validation_start is not None
                and prof.validation_start > 0
            ):
                add_patch(
                    ax,
                    prof.validation_start,
                    prof.validation_stop,
                    colors[8],
                    labels.pop("validation", None),
                )

            # NEW: Add get observation visualization
            if (
                hasattr(prof, "get_observation_start")
                and prof.get_observation_start is not None
                and prof.get_observation_start > 0
            ):
                add_patch(
                    ax,
                    prof.get_observation_start,
                    prof.get_observation_stop,
                    colors[12],
                    labels.pop("get observation", None),
                )

            try:
                next_step_error = step_info_list[i + 1].obs["last_action_error"]
            except (IndexError, KeyError, TypeError):
                next_step_error = ""

            if next_step_error:
                # add a hollow rectangle for error
                label = labels.pop("action error", None)
                add_patch(ax, prof.env_start, prof.env_stop, "red", label, edge=True)

        if step_info.action is not None:
            # Blue rectangle for agent_start to agent_stop
            add_patch(ax, prof.agent_start, prof.agent_stop, colors[10], labels.pop("agent", None))

            # Black vertical bar at agent stop
            ax.axvline(prof.agent_stop, color="black", linewidth=3)
            step_times.append(prof.agent_stop)

            ax.text(
                prof.agent_stop,
                0,
                str(step + 1),
                color="white",
                fontsize=12,
                verticalalignment="bottom",
                horizontalalignment="left",
                rotation=0,
                clip_on=True,
                fontweight=1000,
                backgroundcolor=colors[12],
            )

        if step_info.truncated or step_info.terminated:
            if step_info.truncated:
                color = "black"
            elif step_info.terminated:
                if summary_info.get("cum_reward", 0) > 0:
                    color = "limegreen"
                else:
                    color = "black"

            ax.axvline(prof.env_stop, color=color, linewidth=4, linestyle=":")

            text = f"R:{summary_info.get('cum_reward', np.nan):.1f}"

            if summary_info["err_msg"]:
                text = "Err"
                color = "red"

            ax.text(
                prof.env_stop,
                0.98,
                text,
                color="white",
                fontsize=12,
                verticalalignment="top",
                horizontalalignment="right",
                rotation=0,
                clip_on=True,
                # antialiased=True,
                fontweight=1000,
                backgroundcolor=color,
            )

    ax.set_ylim(0, 1)
    ax.set_xlim(0, max(all_times) + 1)

    ax.set_xlabel("Time")
    ax.set_yticks([])

    # position legend above outside the fig in one row
    ax.legend(
        loc="upper center",
        bbox_to_anchor=(0.5, 1.2),
        ncol=8,  # Updated to accommodate new labels
        frameon=True,
    )

    return step_times

def main():
    run_gradio(RESULTS_DIR)

if __name__ == "__main__":
    main()


main

                backgroundcolor=color,
            )

    ax.set_ylim(0, 1)
    ax.set_xlim(0, max(all_times) + 1)

    ax.set_xlabel("Time")
    ax.set_yticks([])

    # position legend above outside the fig in one row
    ax.legend(
        loc="upper center",
        bbox_to_anchor=(0.5, 1.2),
        ncol=8,  # Updated to accommodate new labels
        frameon=True,
    )

    return step_times

def main():
    run_gradio(RESULTS_DIR)

if __name__ == "__main__":
    main()


    df = df.copy()
    df.columns = clean_column_names(df.columns)
    df.index.names = clean_column_names(df.index.names)
    return df

def remove_args_from_col(df: pd.DataFrame):
    df.columns = [col.replace("_args", "") for col in df.columns]
    df.index.names = [col.replace("_args", "") for col in df.index.names]
    return df

def clean_column_names(col_list):
    # col_list = [col.replace("_args", "") for col in col_list]
    col_list = [col.replace(".", ".\n") for col in col_list]  # adding space for word wrap
    # col_list = [col.replace("_", " ") for col in col_list]
    return col_list

class ClickMapper:
    def __init__(self, ax: plt.Axes, step_times: list[float]):
        self.ax = ax
        self.step_times = step_times

    def to_time(self, x_pix_coord):
        x_time_coord, _ = self.ax.transData.inverted().transform((x_pix_coord, 0))
        return x_time_coord

    def to_step(self, x_pix_coord):
        x_time_coord = self.to_time(x_pix_coord)
        return np.searchsorted(self.step_times, x_time_coord)

@dataclass
class EpisodeId:
    agent_id: str = None
    task_name: str = None
    seed: int = None
    row_index: int = None  # unique row index to disambiguate selections

@dataclass
class StepId:


def remove_args_from_col(df: pd.DataFrame):
    df.columns = [col.replace("_args", "") for col in df.columns]
    df.index.names = [col.replace("_args", "") for col in df.index.names]
    return df

def clean_column_names(col_list):
    # col_list = [col.replace("_args", "") for col in col_list]
    col_list = [col.replace(".", ".\n") for col in col_list]  # adding space for word wrap
    # col_list = [col.replace("_", " ") for col in col_list]
    return col_list

class ClickMapper:
    def __init__(self, ax: plt.Axes, step_times: list[float]):
        self.ax = ax
        self.step_times = step_times

    def to_time(self, x_pix_coord):
        x_time_coord, _ = self.ax.transData.inverted().transform((x_pix_coord, 0))
        return x_time_coord

    def to_step(self, x_pix_coord):
        x_time_coord = self.to_time(x_pix_coord)
        return np.searchsorted(self.step_times, x_time_coord)

@dataclass
class EpisodeId:
    agent_id: str = None
    task_name: str = None
    seed: int = None
    row_index: int = None  # unique row index to disambiguate selections

@dataclass
class StepId:
    episode_id: EpisodeId = None
    step: int = None


    df.index.names = [col.replace("_args", "") for col in df.index.names]
    return df

def clean_column_names(col_list):
    # col_list = [col.replace("_args", "") for col in col_list]
    col_list = [col.replace(".", ".\n") for col in col_list]  # adding space for word wrap
    # col_list = [col.replace("_", " ") for col in col_list]
    return col_list

class ClickMapper:
    def __init__(self, ax: plt.Axes, step_times: list[float]):
        self.ax = ax
        self.step_times = step_times

    def to_time(self, x_pix_coord):
        x_time_coord, _ = self.ax.transData.inverted().transform((x_pix_coord, 0))
        return x_time_coord

    def to_step(self, x_pix_coord):
        x_time_coord = self.to_time(x_pix_coord)
        return np.searchsorted(self.step_times, x_time_coord)

@dataclass
class EpisodeId:
    agent_id: str = None
    task_name: str = None
    seed: int = None
    row_index: int = None  # unique row index to disambiguate selections

@dataclass
class StepId:
    episode_id: EpisodeId = None
    step: int = None

@dataclass
class Info:
    results_dir: Path = None  # to root directory of all experiments
    study_dirs: Path = None  # the path of the currently selected experiment


@dataclass
class StepId:
    episode_id: EpisodeId = None
    step: int = None

@dataclass
class Info:
    results_dir: Path = None  # to root directory of all experiments
    study_dirs: Path = None  # the path of the currently selected experiment
    result_df: pd.DataFrame = None  # the raw loaded df
    agent_df: pd.DataFrame = None  # the df filtered for selected agent
    tasks_df: pd.DataFrame = None  # the unique tasks for selected agent
    exp_result: ExpResult = None  # the selected episode
    click_mapper: ClickMapper = None  # mapping from profiler click to step
    step: int = None  # currently selected step
    active_tab: str = "Screenshot"  # currently selected observation tab
    agent_id_keys: list[str] = None  # the list of columns identifying an agent

    def update_exp_result(self, episode_id: EpisodeId):
        if self.result_df is None or episode_id.task_name is None or episode_id.seed is None:
            self.exp_result = None

        # find unique row using idx
        result_df = self.agent_df.reset_index(inplace=False)
        sub_df = result_df[result_df["_row_index"] == episode_id.row_index]
        if len(sub_df) == 0:
            self.exp_result = None
            raise ValueError(f"Could not find _row_index: {episode_id.row_index}")

        if len(sub_df) > 1:
            warning(
                f"Found multiple rows with same row_index {episode_id.row_index} Using the first one."
            )
        exp_dir = sub_df.iloc[0]["exp_dir"]
        print(exp_dir)
        self.exp_result = ExpResult(exp_dir)
        self.step = 0

    def get_agent_id(self, row: pd.Series):
        agent_id = []
        for key in self.agent_id_keys:
            agent_id.append((key, row[key]))
        return agent_id

    def filter_agent_id(self, agent_id: list[tuple]):
        # Preserve a stable row index to disambiguate selections later
        tmp_df = self.result_df.reset_index(inplace=False)
        tmp_df["_row_index"] = tmp_df.index
        tmp_df.set_index(TASK_NAME_KEY, inplace=True)

        for col, val in agent_id:
            col = col.replace(".\n", ".")
            tmp_df = tmp_df[tmp_df[col] == val]
        self.agent_df = tmp_df

info = Info()


    def update_exp_result(self, episode_id: EpisodeId):
        if self.result_df is None or episode_id.task_name is None or episode_id.seed is None:
            self.exp_result = None

        # find unique row using idx
        result_df = self.agent_df.reset_index(inplace=False)
        sub_df = result_df[result_df["_row_index"] == episode_id.row_index]
        if len(sub_df) == 0:
            self.exp_result = None
            raise ValueError(f"Could not find _row_index: {episode_id.row_index}")

        if len(sub_df) > 1:
            warning(
                f"Found multiple rows with same row_index {episode_id.row_index} Using the first one."
            )
        exp_dir = sub_df.iloc[0]["exp_dir"]
        print(exp_dir)
        self.exp_result = ExpResult(exp_dir)
        self.step = 0

    def get_agent_id(self, row: pd.Series):
        agent_id = []
        for key in self.agent_id_keys:
            agent_id.append((key, row[key]))
        return agent_id

    def filter_agent_id(self, agent_id: list[tuple]):
        # Preserve a stable row index to disambiguate selections later
        tmp_df = self.result_df.reset_index(inplace=False)
        tmp_df["_row_index"] = tmp_df.index
        tmp_df.set_index(TASK_NAME_KEY, inplace=True)

        for col, val in agent_id:
            col = col.replace(".\n", ".")
            tmp_df = tmp_df[tmp_df[col] == val]
        self.agent_df = tmp_df

info = Info()

css = """
.my-markdown {
    max-height: 400px;
    overflow-y: auto;


        sub_df = result_df[result_df["_row_index"] == episode_id.row_index]
        if len(sub_df) == 0:
            self.exp_result = None
            raise ValueError(f"Could not find _row_index: {episode_id.row_index}")

        if len(sub_df) > 1:
            warning(
                f"Found multiple rows with same row_index {episode_id.row_index} Using the first one."
            )
        exp_dir = sub_df.iloc[0]["exp_dir"]
        print(exp_dir)
        self.exp_result = ExpResult(exp_dir)
        self.step = 0

    def get_agent_id(self, row: pd.Series):
        agent_id = []
        for key in self.agent_id_keys:
            agent_id.append((key, row[key]))
        return agent_id

    def filter_agent_id(self, agent_id: list[tuple]):
        # Preserve a stable row index to disambiguate selections later
        tmp_df = self.result_df.reset_index(inplace=False)
        tmp_df["_row_index"] = tmp_df.index
        tmp_df.set_index(TASK_NAME_KEY, inplace=True)

        for col, val in agent_id:
            col = col.replace(".\n", ".")
            tmp_df = tmp_df[tmp_df[col] == val]
        self.agent_df = tmp_df

info = Info()

css = """
.my-markdown {
    max-height: 400px;
    overflow-y: auto;
}
.error-report {
    max-height: 700px;
    overflow-y: auto;
}
.my-code-view {
    max-height: 300px;
    overflow-y: auto;
}
code {
    white-space: pre-wrap;


decorator

        # print(f"Key event: {key_event}")
        step = step_id.step
        if key_event.startswith("Cmd+Left"):
            step = max(0, step - 1)
        elif key_event.startswith("Cmd+Right"):
            step = min(len(info.exp_result.steps_info) - 2, step + 1)
        else:
            return gr.update()
        # print(f"Updating step to {step} from key event {key_event}")
        info.step = step
        step_id = StepId(episode_id=step_id.episode_id, step=step)
    return ("", step_id)

def tab_select(evt: gr.SelectData):
    global info
    info.active_tab = evt.value

def if_active(tab_name, n_out=1):
    def decorator(fn):
        def wrapper(*args, **kwargs):
            global info
            if info.active_tab == tab_name:
                # print("updating: ", fn.__name__)
                return fn(*args, **kwargs)
            else:
                # print("skipping: ", fn.__name__)
                if n_out == 1:
                    return gr.update()
                elif n_out > 1:
                    return (gr.update(),) * n_out

        return wrapper

    return decorator

def update_screenshot(som_or_not: str):
    global info
    img, action_str = get_screenshot(info, som_or_not=som_or_not, annotate=True)
    return img

def get_screenshot(
    info: Info, step: int = None, som_or_not: str = "Raw Screenshots", annotate: bool = False
):
    if step is None:
        step = info.step
    try:
        step_info = info.exp_result.steps_info[step]
        is_som = som_or_not == "SOM Screenshots"
        img = info.exp_result.get_screenshot(step, som=is_som)
        if annotate:


def get_state_error(state: Info):
    try:
        step_info = state.exp_result.steps_info[state.step + 1]
        err_msg = step_info.obs.get("last_action_error", None)
    except (IndexError, AttributeError):
        err_msg = None

    if err_msg is None or len(err_msg) == 0:
        err_msg = "No Error"
    return f"""\
**Step error after action:**

{code(err_msg)}"""

def get_seeds_df(result_df: pd.DataFrame, task_name: str):
    result_df = result_df.reset_index(inplace=False)
    result_df = result_df[result_df[TASK_NAME_KEY] == task_name]

    def extract_columns(row: pd.Series):
        return pd.Series(
            {
                "idx": row.get("_row_index", None),
                "seed": row.get(TASK_SEED_KEY, None),
                "reward": row.get("cum_reward", None),
                "err": bool(row.get("err_msg", None)),
                "n_steps": row.get("n_steps", None),
            }
        )

    seed_df = result_df.apply(extract_columns, axis=1)
    # Ensure column order and readability
    seed_df = seed_df[["seed", "reward", "err", "n_steps", "idx"]]
    return seed_df

def on_select_agent(evt: gr.SelectData, df: pd.DataFrame):
    # TODO try to find a clever way to solve the sort bug here
    return info.get_agent_id(df.iloc[evt.index[0]])

def on_select_task(evt: gr.SelectData, df: pd.DataFrame, agent_id: list[tuple]):
    # get col index
    col_idx = df.columns.get_loc(TASK_NAME_KEY)
    return (agent_id, evt.row_value[col_idx])

def update_seeds(agent_task_id: tuple):
    agent_id, task_name = agent_task_id


wrapper

        step = step_id.step
        if key_event.startswith("Cmd+Left"):
            step = max(0, step - 1)
        elif key_event.startswith("Cmd+Right"):
            step = min(len(info.exp_result.steps_info) - 2, step + 1)
        else:
            return gr.update()
        # print(f"Updating step to {step} from key event {key_event}")
        info.step = step
        step_id = StepId(episode_id=step_id.episode_id, step=step)
    return ("", step_id)

def tab_select(evt: gr.SelectData):
    global info
    info.active_tab = evt.value

def if_active(tab_name, n_out=1):
    def decorator(fn):
        def wrapper(*args, **kwargs):
            global info
            if info.active_tab == tab_name:
                # print("updating: ", fn.__name__)
                return fn(*args, **kwargs)
            else:
                # print("skipping: ", fn.__name__)
                if n_out == 1:
                    return gr.update()
                elif n_out > 1:
                    return (gr.update(),) * n_out

        return wrapper

    return decorator

def update_screenshot(som_or_not: str):
    global info
    img, action_str = get_screenshot(info, som_or_not=som_or_not, annotate=True)
    return img

def get_screenshot(
    info: Info, step: int = None, som_or_not: str = "Raw Screenshots", annotate: bool = False
):
    if step is None:
        step = info.step
    try:
        step_info = info.exp_result.steps_info[step]
        is_som = som_or_not == "SOM Screenshots"


import logging
import os
import time
from functools import partial
from typing import Any, List, Optional, Union

from pydantic import Field

from agentlab.llm.base_api import AbstractChatModel
from agentlab.llm.llm_utils import AIMessage, Discussion
from agentlab.llm.prompt_templates import PromptTemplate, get_prompt_template

class HFBaseChatModel(AbstractChatModel):
    """
    Custom LLM Chatbot that can interface with HuggingFace models with support for multiple samples.

    This class allows for the creation of a custom chatbot using models hosted
    on HuggingFace Hub or a local checkpoint. It provides flexibility in defining
    the temperature for response sampling and the maximum number of new tokens
    in the response.

    Attributes:
        llm (Any): The HuggingFaceHub model instance.
        prompt_template (Any): Template for the prompt to be used for the model's input sequence.
        tokenizer (Any): The tokenizer to use for the model.
        n_retry_server (int): Number of times to retry on server failure.
    """

    llm: Any = Field(description="The HuggingFaceHub model instance")
    tokenizer: Any = Field(
        default=None,
        description="The tokenizer to use for the model",
    )
    prompt_template: Optional[PromptTemplate] = Field(
        default=None,
        description="Template for the prompt to be used for the model's input sequence",
    )
    n_retry_server: int = Field(
        default=4,
        description="The number of times to retry the server if it fails to respond",
    )

    def __init__(self, model_name, base_model_name, n_retry_server, log_probs):
        super().__init__()
        self.n_retry_server = n_retry_server
        self.log_probs = log_probs

        # Lazy import to avoid heavy transformers import when unused
        try:
            from transformers import AutoTokenizer, GPT2TokenizerFast  # type: ignore
        except Exception as e:  # pragma: no cover - surfaced only when transformers missing
            raise ImportError(
                "The 'transformers' package is required for HuggingFace models. Install it to use HF backends."
            ) from e

        if base_model_name is None:
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        else:
            self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
        if isinstance(self.tokenizer, GPT2TokenizerFast):
            logging.warning(
                f"No chat template is defined for {base_model_name}. Resolving to the hard-coded templates."
            )
            self.tokenizer = None
            self.prompt_template = get_prompt_template(model_name)

    def __call__(
        self,
        messages: list[dict],
        n_samples: int = 1,
        temperature: Optional[float] = None,
    ) -> Union[AIMessage, List[AIMessage]]:
        """
        Generate one or more responses for the given messages.

        Args:
            messages: List of message dictionaries containing the conversation history.
            n_samples: Number of independent responses to generate. Defaults to 1.
            temperature: The temperature for response sampling. Defaults to None.

        Returns:
            If n_samples=1, returns a single AIMessage.
            If n_samples>1, returns a list of AIMessages.

        Raises:
            Exception: If the server fails to respond after n_retry_server attempts or if the chat template fails.
        """
        if self.tokenizer:
            try:
                if isinstance(messages, Discussion):
                    messages.merge()
                prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
            except Exception as e:
                if "Conversation roles must alternate" in str(e):
                    logging.warning(
                        "Failed to apply the chat template. Maybe because it doesn't support the 'system' role. "
                        "Retrying with the 'system' role appended to the 'user' role."
                    )
                    messages = _prepend_system_to_first_user(messages)
                    prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
                else:
                    raise e
        elif self.prompt_template:
            prompt = self.prompt_template.construct_prompt(messages)

        responses = []
        for _ in range(n_samples):
            itr = 0
            while True:
                try:
                    temperature = (
                        temperature
                        if temperature is not None
                        else getattr(self, "temperature", 0.1)
                    )
                    answer = self.llm(prompt, temperature=temperature)
                    response = AIMessage(answer)
                    if self.log_probs:
                        response["content"] = answer.generated_text
                        response["log_probs"] = answer.details
                    responses.append(response)
                    break
                except Exception as e:
                    if itr == self.n_retry_server - 1:
                        raise e
                    logging.warning(
                        f"Failed to get a response from the server: \n{e}\n"
                        f"Retrying... ({itr+1}/{self.n_retry_server})"
                    )
                    time.sleep(5)
                    itr += 1

        return responses[0] if n_samples == 1 else responses

    def _llm_type(self):
        return "huggingface"

def _prepend_system_to_first_user(messages, column_remap={}):
    # Initialize an index for the system message
    system_index = None

    human_key = column_remap.get("HumanMessage", "user")
    role_key = column_remap.get("role", "role")
    text_key = column_remap.get("text", "content")

    # Find the system content and its index
    for i, msg in enumerate(messages):
        if msg[role_key] == "system":
            system_index = i
            system_content = msg[text_key]
            break  # Stop after finding the first system message

    # If a system message was found, modify the first user message and remove the system message
    if system_index is not None:
        for msg in messages:
            if msg[role_key] == human_key:
                # Prepend system content to the first user content
                msg[text_key] = str(system_content) + "\n" + str(msg[text_key])
                # Remove the original system message
                del messages[system_index]
                break  # Ensures that only the first user message is modified

    return messages

class HuggingFaceURLChatModel(HFBaseChatModel):
    """HF backend using a Text Generation Inference (TGI) HTTP endpoint.

    This class is placed here to keep all heavy HF imports optional and only
    loaded when a HF backend is explicitly requested.
    """

    def __init__(
        self,
        model_name: str,
        model_url: str,
        base_model_name: Optional[str] = None,
        token: Optional[str] = None,
        temperature: Optional[float] = 1e-1,
        max_new_tokens: Optional[int] = 512,
        n_retry_server: Optional[int] = 4,
        log_probs: Optional[bool] = False,
    ):
        super().__init__(model_name, base_model_name, n_retry_server, log_probs)
        if temperature is not None and temperature < 1e-3:
            logging.warning("Models might behave weirdly when temperature is too low.")
        self.temperature = temperature

        if token is None:
            # support both env var names used elsewhere
            token = os.environ.get("TGI_TOKEN") or os.environ.get("AGENTLAB_MODEL_TOKEN")

        # Lazy import huggingface_hub here to avoid import on non-HF paths
        try:
            from huggingface_hub import InferenceClient  # type: ignore
        except Exception as e:  # pragma: no cover - surfaced only when package missing
            raise ImportError(
                "The 'huggingface_hub' package is required for HuggingFace URL backends."
            ) from e

        client = InferenceClient(model=model_url, token=token)
        self.llm = partial(
            client.text_generation,
            max_new_tokens=max_new_tokens,
            details=log_probs,
        )


HFBaseChatModel

import logging
import os
import time
from functools import partial
from typing import Any, List, Optional, Union

from pydantic import Field

from agentlab.llm.base_api import AbstractChatModel
from agentlab.llm.llm_utils import AIMessage, Discussion
from agentlab.llm.prompt_templates import PromptTemplate, get_prompt_template

class HFBaseChatModel(AbstractChatModel):
    """
    Custom LLM Chatbot that can interface with HuggingFace models with support for multiple samples.

    This class allows for the creation of a custom chatbot using models hosted
    on HuggingFace Hub or a local checkpoint. It provides flexibility in defining
    the temperature for response sampling and the maximum number of new tokens
    in the response.

    Attributes:
        llm (Any): The HuggingFaceHub model instance.
        prompt_template (Any): Template for the prompt to be used for the model's input sequence.
        tokenizer (Any): The tokenizer to use for the model.
        n_retry_server (int): Number of times to retry on server failure.
    """

    llm: Any = Field(description="The HuggingFaceHub model instance")
    tokenizer: Any = Field(
        default=None,
        description="The tokenizer to use for the model",
    )
    prompt_template: Optional[PromptTemplate] = Field(
        default=None,
        description="Template for the prompt to be used for the model's input sequence",
    )
    n_retry_server: int = Field(
        default=4,
        description="The number of times to retry the server if it fails to respond",
    )

    def __init__(self, model_name, base_model_name, n_retry_server, log_probs):
        super().__init__()
        self.n_retry_server = n_retry_server
        self.log_probs = log_probs

        # Lazy import to avoid heavy transformers import when unused
        try:
            from transformers import AutoTokenizer, GPT2TokenizerFast  # type: ignore
        except Exception as e:  # pragma: no cover - surfaced only when transformers missing
            raise ImportError(
                "The 'transformers' package is required for HuggingFace models. Install it to use HF backends."
            ) from e

        if base_model_name is None:
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        else:
            self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
        if isinstance(self.tokenizer, GPT2TokenizerFast):
            logging.warning(
                f"No chat template is defined for {base_model_name}. Resolving to the hard-coded templates."
            )
            self.tokenizer = None
            self.prompt_template = get_prompt_template(model_name)

    def __call__(
        self,
        messages: list[dict],
        n_samples: int = 1,
        temperature: Optional[float] = None,
    ) -> Union[AIMessage, List[AIMessage]]:
        """
        Generate one or more responses for the given messages.

        Args:
            messages: List of message dictionaries containing the conversation history.
            n_samples: Number of independent responses to generate. Defaults to 1.
            temperature: The temperature for response sampling. Defaults to None.

        Returns:
            If n_samples=1, returns a single AIMessage.
            If n_samples>1, returns a list of AIMessages.

        Raises:
            Exception: If the server fails to respond after n_retry_server attempts or if the chat template fails.
        """
        if self.tokenizer:
            try:
                if isinstance(messages, Discussion):
                    messages.merge()
                prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
            except Exception as e:
                if "Conversation roles must alternate" in str(e):
                    logging.warning(
                        "Failed to apply the chat template. Maybe because it doesn't support the 'system' role. "
                        "Retrying with the 'system' role appended to the 'user' role."
                    )
                    messages = _prepend_system_to_first_user(messages)
                    prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
                else:
                    raise e
        elif self.prompt_template:
            prompt = self.prompt_template.construct_prompt(messages)

        responses = []
        for _ in range(n_samples):
            itr = 0
            while True:
                try:
                    temperature = (
                        temperature
                        if temperature is not None
                        else getattr(self, "temperature", 0.1)
                    )
                    answer = self.llm(prompt, temperature=temperature)
                    response = AIMessage(answer)
                    if self.log_probs:
                        response["content"] = answer.generated_text
                        response["log_probs"] = answer.details
                    responses.append(response)
                    break
                except Exception as e:
                    if itr == self.n_retry_server - 1:
                        raise e
                    logging.warning(
                        f"Failed to get a response from the server: \n{e}\n"
                        f"Retrying... ({itr+1}/{self.n_retry_server})"
                    )
                    time.sleep(5)
                    itr += 1

        return responses[0] if n_samples == 1 else responses

    def _llm_type(self):
        return "huggingface"

def _prepend_system_to_first_user(messages, column_remap={}):
    # Initialize an index for the system message
    system_index = None

    human_key = column_remap.get("HumanMessage", "user")
    role_key = column_remap.get("role", "role")
    text_key = column_remap.get("text", "content")

    # Find the system content and its index
    for i, msg in enumerate(messages):
        if msg[role_key] == "system":
            system_index = i
            system_content = msg[text_key]
            break  # Stop after finding the first system message

    # If a system message was found, modify the first user message and remove the system message
    if system_index is not None:
        for msg in messages:


                        response["content"] = answer.generated_text
                        response["log_probs"] = answer.details
                    responses.append(response)
                    break
                except Exception as e:
                    if itr == self.n_retry_server - 1:
                        raise e
                    logging.warning(
                        f"Failed to get a response from the server: \n{e}\n"
                        f"Retrying... ({itr+1}/{self.n_retry_server})"
                    )
                    time.sleep(5)
                    itr += 1

        return responses[0] if n_samples == 1 else responses

    def _llm_type(self):
        return "huggingface"

def _prepend_system_to_first_user(messages, column_remap={}):
    # Initialize an index for the system message
    system_index = None

    human_key = column_remap.get("HumanMessage", "user")
    role_key = column_remap.get("role", "role")
    text_key = column_remap.get("text", "content")

    # Find the system content and its index
    for i, msg in enumerate(messages):
        if msg[role_key] == "system":
            system_index = i
            system_content = msg[text_key]
            break  # Stop after finding the first system message

    # If a system message was found, modify the first user message and remove the system message
    if system_index is not None:
        for msg in messages:
            if msg[role_key] == human_key:
                # Prepend system content to the first user content
                msg[text_key] = str(system_content) + "\n" + str(msg[text_key])
                # Remove the original system message
                del messages[system_index]
                break  # Ensures that only the first user message is modified

    return messages

class HuggingFaceURLChatModel(HFBaseChatModel):
    """HF backend using a Text Generation Inference (TGI) HTTP endpoint.

    This class is placed here to keep all heavy HF imports optional and only
    loaded when a HF backend is explicitly requested.
    """

    def __init__(
        self,
        model_name: str,
        model_url: str,
        base_model_name: Optional[str] = None,
        token: Optional[str] = None,
        temperature: Optional[float] = 1e-1,
        max_new_tokens: Optional[int] = 512,
        n_retry_server: Optional[int] = 4,
        log_probs: Optional[bool] = False,
    ):


HuggingFaceURLChatModel

    # Find the system content and its index
    for i, msg in enumerate(messages):
        if msg[role_key] == "system":
            system_index = i
            system_content = msg[text_key]
            break  # Stop after finding the first system message

    # If a system message was found, modify the first user message and remove the system message
    if system_index is not None:
        for msg in messages:
            if msg[role_key] == human_key:
                # Prepend system content to the first user content
                msg[text_key] = str(system_content) + "\n" + str(msg[text_key])
                # Remove the original system message
                del messages[system_index]
                break  # Ensures that only the first user message is modified

    return messages

class HuggingFaceURLChatModel(HFBaseChatModel):
    """HF backend using a Text Generation Inference (TGI) HTTP endpoint.

    This class is placed here to keep all heavy HF imports optional and only
    loaded when a HF backend is explicitly requested.
    """

    def __init__(
        self,
        model_name: str,
        model_url: str,
        base_model_name: Optional[str] = None,
        token: Optional[str] = None,
        temperature: Optional[float] = 1e-1,
        max_new_tokens: Optional[int] = 512,
        n_retry_server: Optional[int] = 4,
        log_probs: Optional[bool] = False,
    ):
        super().__init__(model_name, base_model_name, n_retry_server, log_probs)
        if temperature is not None and temperature < 1e-3:
            logging.warning("Models might behave weirdly when temperature is too low.")
        self.temperature = temperature

        if token is None:
            # support both env var names used elsewhere
            token = os.environ.get("TGI_TOKEN") or os.environ.get("AGENTLAB_MODEL_TOKEN")

        # Lazy import huggingface_hub here to avoid import on non-HF paths
        try:
            from huggingface_hub import InferenceClient  # type: ignore
        except Exception as e:  # pragma: no cover - surfaced only when package missing
            raise ImportError(
                "The 'huggingface_hub' package is required for HuggingFace URL backends."
            ) from e

        client = InferenceClient(model=model_url, token=token)
        self.llm = partial(
            client.text_generation,
            max_new_tokens=max_new_tokens,
            details=log_probs,
        )


    # If a system message was found, modify the first user message and remove the system message
    if system_index is not None:
        for msg in messages:
            if msg[role_key] == human_key:
                # Prepend system content to the first user content
                msg[text_key] = str(system_content) + "\n" + str(msg[text_key])
                # Remove the original system message
                del messages[system_index]
                break  # Ensures that only the first user message is modified

    return messages

class HuggingFaceURLChatModel(HFBaseChatModel):
    """HF backend using a Text Generation Inference (TGI) HTTP endpoint.

    This class is placed here to keep all heavy HF imports optional and only
    loaded when a HF backend is explicitly requested.
    """

    def __init__(
        self,
        model_name: str,
        model_url: str,
        base_model_name: Optional[str] = None,
        token: Optional[str] = None,
        temperature: Optional[float] = 1e-1,
        max_new_tokens: Optional[int] = 512,
        n_retry_server: Optional[int] = 4,
        log_probs: Optional[bool] = False,
    ):
        super().__init__(model_name, base_model_name, n_retry_server, log_probs)
        if temperature is not None and temperature < 1e-3:
            logging.warning("Models might behave weirdly when temperature is too low.")
        self.temperature = temperature

        if token is None:
            # support both env var names used elsewhere
            token = os.environ.get("TGI_TOKEN") or os.environ.get("AGENTLAB_MODEL_TOKEN")

        # Lazy import huggingface_hub here to avoid import on non-HF paths
        try:
            from huggingface_hub import InferenceClient  # type: ignore
        except Exception as e:  # pragma: no cover - surfaced only when package missing
            raise ImportError(
                "The 'huggingface_hub' package is required for HuggingFace URL backends."
            ) from e

        client = InferenceClient(model=model_url, token=token)
        self.llm = partial(
            client.text_generation,
            max_new_tokens=max_new_tokens,
            details=log_probs,
        )


        # Lazy import to avoid heavy transformers import when unused
        try:
            from transformers import AutoTokenizer, GPT2TokenizerFast  # type: ignore
        except Exception as e:  # pragma: no cover - surfaced only when transformers missing
            raise ImportError(
                "The 'transformers' package is required for HuggingFace models. Install it to use HF backends."
            ) from e

        if base_model_name is None:
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        else:
            self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
        if isinstance(self.tokenizer, GPT2TokenizerFast):
            logging.warning(
                f"No chat template is defined for {base_model_name}. Resolving to the hard-coded templates."
            )
            self.tokenizer = None
            self.prompt_template = get_prompt_template(model_name)

    def __call__(
        self,
        messages: list[dict],
        n_samples: int = 1,
        temperature: Optional[float] = None,
    ) -> Union[AIMessage, List[AIMessage]]:
        """
        Generate one or more responses for the given messages.

        Args:
            messages: List of message dictionaries containing the conversation history.
            n_samples: Number of independent responses to generate. Defaults to 1.
            temperature: The temperature for response sampling. Defaults to None.

        Returns:
            If n_samples=1, returns a single AIMessage.
            If n_samples>1, returns a list of AIMessages.

        Raises:
            Exception: If the server fails to respond after n_retry_server attempts or if the chat template fails.
        """
        if self.tokenizer:
            try:
                if isinstance(messages, Discussion):
                    messages.merge()
                prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
            except Exception as e:
                if "Conversation roles must alternate" in str(e):
                    logging.warning(
                        "Failed to apply the chat template. Maybe because it doesn't support the 'system' role. "
                        "Retrying with the 'system' role appended to the 'user' role."
                    )
                    messages = _prepend_system_to_first_user(messages)
                    prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
                else:
                    raise e
        elif self.prompt_template:
            prompt = self.prompt_template.construct_prompt(messages)

        responses = []
        for _ in range(n_samples):
            itr = 0
            while True:
                try:
                    temperature = (
                        temperature
                        if temperature is not None
                        else getattr(self, "temperature", 0.1)
                    )
                    answer = self.llm(prompt, temperature=temperature)
                    response = AIMessage(answer)
                    if self.log_probs:
                        response["content"] = answer.generated_text
                        response["log_probs"] = answer.details
                    responses.append(response)
                    break
                except Exception as e:
                    if itr == self.n_retry_server - 1:
                        raise e
                    logging.warning(
                        f"Failed to get a response from the server: \n{e}\n"
                        f"Retrying... ({itr+1}/{self.n_retry_server})"
                    )
                    time.sleep(5)
                    itr += 1

        return responses[0] if n_samples == 1 else responses

    def _llm_type(self):
        return "huggingface"

def _prepend_system_to_first_user(messages, column_remap={}):
    # Initialize an index for the system message
    system_index = None

    human_key = column_remap.get("HumanMessage", "user")
    role_key = column_remap.get("role", "role")
    text_key = column_remap.get("text", "content")

    # Find the system content and its index
    for i, msg in enumerate(messages):
        if msg[role_key] == "system":
            system_index = i
            system_content = msg[text_key]
            break  # Stop after finding the first system message


                    )
                    answer = self.llm(prompt, temperature=temperature)
                    response = AIMessage(answer)
                    if self.log_probs:
                        response["content"] = answer.generated_text
                        response["log_probs"] = answer.details
                    responses.append(response)
                    break
                except Exception as e:
                    if itr == self.n_retry_server - 1:
                        raise e
                    logging.warning(
                        f"Failed to get a response from the server: \n{e}\n"
                        f"Retrying... ({itr+1}/{self.n_retry_server})"
                    )
                    time.sleep(5)
                    itr += 1

        return responses[0] if n_samples == 1 else responses

    def _llm_type(self):
        return "huggingface"

def _prepend_system_to_first_user(messages, column_remap={}):
    # Initialize an index for the system message
    system_index = None

    human_key = column_remap.get("HumanMessage", "user")
    role_key = column_remap.get("role", "role")
    text_key = column_remap.get("text", "content")

    # Find the system content and its index
    for i, msg in enumerate(messages):
        if msg[role_key] == "system":
            system_index = i
            system_content = msg[text_key]
            break  # Stop after finding the first system message

    # If a system message was found, modify the first user message and remove the system message
    if system_index is not None:
        for msg in messages:


import json
import logging
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Literal, Optional, Union

import openai
from anthropic import Anthropic
from anthropic.types import Completion
from anthropic.types import Message as AnthrophicMessage
from openai import OpenAI

from agentlab.llm.llm_utils import image_to_png_base64_url

from .base_api import BaseModelArgs
from .llm_utils import (
    call_anthropic_api_with_retries,
    call_openai_api_with_retries,
)
from .tracking import TrackAPIPricingMixin

"""This module contains utlity classes for building input messages and interacting with LLM APIs.
It includes:
    1. Message Builder for building input messages
    2. Base Reponse class for different LLM APIs (OpenAI, Anthropic, etc.)
    3. Factory classes (inherits from BaseModelArgs) for creating instances of LLM Response models.
"""

logger = logging.getLogger(__name__)

ContentItem = Dict[str, Any]
Message = Dict[str, Union[str, List[ContentItem]]]

@dataclass
class ToolCall:
    """Represents a tool call made by the LLM.
    Attributes:
    name: Name of the tool called.
    arguments: Arguments passed to the tool.
    raw_call: The raw call object from the LLM API.
    tool_response: Output of the tool call goes here. It can be only one content item.
    """

    name: str = field(default=None)
    arguments: Dict[str, Any] = field(default_factory=dict)
    raw_call: Any = field(default=None)
    tool_response: ContentItem = None

    @property
    def is_response_set(self) -> bool:
        """Check if the tool response is set."""
        return self.tool_response is not None

    def response_text(self, text: str) -> "MessageBuilder":
        self.tool_response = {"text": text}
        return self

    def response_image(self, image: str) -> "MessageBuilder":
        self.tool_response = {"image": image}
        return self

    def __repr__(self):
        return f"ToolCall(name={self.name}, arguments={self.arguments})"

@dataclass
class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self

    @property
    def all_responses_set(self) -> bool:
        """Check if all tool calls have responses set."""
        return all(call.is_response_set for call in self.tool_calls)

    def __len__(self) -> int:
        """Return the number of tool calls."""
        return len(self.tool_calls)

    def __iter__(self):
        """Make ToolCalls iterable."""
        return iter(self.tool_calls)

    def __bool__(self):
        """Check if there are any tool calls."""
        return len(self.tool_calls) > 0

@dataclass
class LLMOutput:
    """Serializable object for the output of a response LLM."""

    raw_response: Any = field(default=None)
    think: str = field(default="")
    action: str | None = field(default=None)  # Default action if no tool call is made
    tool_calls: ToolCalls | None = field(
        default=None
    )  # This will hold the tool call response if any

class MessageBuilder:
    def __init__(self, role: str):
        self.role = role
        self.content: List[ContentItem] = []
        self.responded_tool_calls: ToolCalls = None

    @classmethod
    def system(cls) -> "MessageBuilder":
        return cls("system")

    @classmethod
    def user(cls) -> "MessageBuilder":
        return cls("user")

    @classmethod
    def assistant(cls) -> "MessageBuilder":
        return cls("assistant")

    @abstractmethod
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the API call."""
        raise NotImplementedError("Subclasses must implement this method.")

    def add_text(self, text: str) -> "MessageBuilder":
        self.content.append({"text": text})
        return self

    def add_image(self, image: str) -> "MessageBuilder":
        self.content.append({"image": image})
        return self

    def to_markdown(self) -> str:
        parts = []
        for item in self.content:
            if "text" in item:
                parts.append(f"\n```\n{item['text']}\n```\n")
            elif "image" in item:
                parts.append(f"![Image]({item['image']})")

        # Tool call markdown repr
        if self.responded_tool_calls is not None:
            for i, tool_call in enumerate(self.responded_tool_calls.tool_calls, 1):
                parts.append(
                    f"\n**Tool Call {i}**: {tool_call_to_python_code(tool_call.name, tool_call.arguments)}"
                )
                response = tool_call.tool_response
                if response is not None:
                    parts.append(f"\n**Tool Response {i}:**")
                    content = (
                        f"```\n{response['text']}\n```"
                        if "text" in response
                        else f"![Tool Response Image]({response['image']})"
                    )
                    parts.append(content)

        markdown = f"### {self.role.capitalize()}\n"
        markdown += "\n".join(parts)

        return markdown

    def add_image_url(self, image_url: str) -> "MessageBuilder":
        """Add an image URL to the message content."""
        self.content.append({"image": image_to_png_base64_url(image_url)})
        return self

    def mark_all_previous_msg_for_caching(self):
        """Insert a cache breakpoint in the message content."""
        # This is a placeholder for future implementation.
        raise NotImplementedError

    @classmethod
    def add_responded_tool_calls(cls, responded_tool_calls: ToolCalls) -> "MessageBuilder":
        """Add tool calls to the message content."""
        assert responded_tool_calls.all_responses_set, "All tool calls must have a response."
        msg = cls("tool")
        msg.responded_tool_calls = responded_tool_calls
        return msg

class OpenAIResponseAPIMessageBuilder(MessageBuilder):
    @classmethod
    def system(cls) -> "OpenAIResponseAPIMessageBuilder":
        # OpenAI Responses API uses 'developer' role for system messages
        return cls("developer")

    def prepare_message(self) -> List[Message]:
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]

        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Convert the content item to the expected format for OpenAI Responses."""
        if "text" in content:
            content_type = "input_text" if self.role != "assistant" else "output_text"
            return {"type": content_type, "text": content["text"]}
        elif "image" in content:
            return {"type": "input_image", "image_url": content["image"]}
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:
            raise ValueError("No tool calls found in responded_tool_calls")

        output = []
        output.extend(self.responded_tool_calls.raw_calls.output)  # this contains response
        for fn_call in self.responded_tool_calls:
            call_type = fn_call.raw_call.type
            call_id = fn_call.raw_call.call_id
            call_response = fn_call.tool_response

            match call_type:
                case "function_call":
                    # image output is not supported in function calls response.
                    assert (
                        "image" not in call_response
                    ), "Image output is not supported in function calls response."
                    fn_call_response = {
                        "type": "function_call_output",
                        "call_id": call_id,
                        "output": self.convert_content_to_expected_format(call_response)["text"],
                    }
                    output.append(fn_call_response)

                case "computer_call":
                    # For computer calls, use only images are expected.
                    assert (
                        "text" not in call_response
                    ), "Text output is not supported in computer calls response."
                    computer_call_output = {
                        "type": "computer_call_output",
                        "call_id": call_id,
                        "output": self.convert_content_to_expected_format(call_response),
                    }
                    output.append(computer_call_output)  # this needs to be a screenshot

        return output

    def mark_all_previous_msg_for_caching(self):
        """Nothing special to do here for openAI. They do not have a notion of cache breakpoints."""
        pass

class AnthropicAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        content = [self.transform_content(item) for item in self.content]
        output = {"role": self.role, "content": content}

        if self.role == "tool":
            return self.handle_tool_call()

        if self.role == "assistant":
            # Strip whitespace from assistant text responses. See anthropic error code 400.
            for c in output["content"]:
                if "text" in c:
                    c["text"] = c["text"].strip()
        return [output]

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:
            raise ValueError("No tool calls found in responded_tool_calls")

        llm_tool_call = {
            "role": "assistant",
            "content": self.responded_tool_calls.raw_calls.content,
        }  # Add the toolcall block
        tool_response = {"role": "user", "content": []}  # Anthropic expects a list of messages
        for call in self.responded_tool_calls:
            assert (
                "image" not in call.tool_response
            ), "Image output is not supported in tool calls response."
            tool_response["content"].append(
                {
                    "type": "tool_result",
                    "tool_use_id": call.raw_call.id,
                    "content": self.transform_content(call.tool_response)[
                        "text"
                    ],  # needs to be str
                }
            )

        return [llm_tool_call, tool_response]

    def transform_content(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by Anthropic API."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content:
            img_str: str = content["image"]
            # make sure to get rid of the image type for anthropic
            # e.g. "data:image/png;base64"
            if img_str.startswith("data:image/png;base64,"):
                img_str = img_str[len("data:image/png;base64,") :]
            return {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png",
                    "data": img_str,
                },
            }
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def mark_all_previous_msg_for_caching(self) -> List[Message]:
        """Insert a cache breakpoint in the message content to mark all previous messages for caching."""
        self._cache_breakpoint = True

class OpenAIChatCompletionAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the OpenAI API."""
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]
        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by OpenAI ChatCompletion."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content:
            return {"type": "image_url", "image_url": {"url": content["image"]}}
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:
            raise ValueError("No tool calls found in responded_tool_calls")
        output = []
        output.append(
            self.responded_tool_calls.raw_calls.choices[0].message
        )  # add raw calls to output
        for fn_call in self.responded_tool_calls:
            raw_call = fn_call.raw_call
            assert (
                "image" not in fn_call.tool_response
            ), "Image output is not supported in function calls response."
            # a function_call_output dict has keys "role", "tool_call_id" and "content"
            tool_call_reponse = {
                "name": raw_call["function"]["name"],  # required with OpenRouter
                "role": "tool",
                "tool_call_id": raw_call["id"],
                "content": self.convert_content_to_expected_format(fn_call.tool_response)["text"],
            }
            output.append(tool_call_reponse)

        return output

    def mark_all_previous_msg_for_caching(self):
        """Nothing special to do here for openAI. They do not have a notion of cache breakpoints."""
        pass

@dataclass
class APIPayload:
    messages: List[MessageBuilder] | None = None
    tools: List[Dict[str, Any]] | None = None
    tool_choice: Literal["none", "auto", "any", "required"] | None = None
    force_call_tool: str | None = (
        None  # Name of the tool to call # If set, will force the LLM to call this tool.
    )
    use_cache_breakpoints: bool = (
        False  # If True, will apply cache breakpoints to the messages. # applicable for Anthropic
    )
    cache_tool_definition: bool = (
        False  # If True, will cache the tool definition in the last message.
    )
    cache_complete_prompt: bool = (
        False  # If True, will cache the complete prompt in the last message.
    )
    reasoning_effort: Literal["low", "medium", "high"] | None = None

    def __post_init__(self):
        if self.tool_choice and self.force_call_tool:
            raise ValueError("tool_choice and force_call_tool are mutually exclusive")
        if self.reasoning_effort is not None:
            logger.info(
                "In agentlab reasoning_effort is used by LiteLLM API only. We will eventually shift to LiteLLM API for all LLMs."
            )

# # Base class for all API Endpoints
class BaseResponseModel(ABC):
    def __init__(
        self,
        model_name: str,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ):
        self.model_name = model_name
        self.api_key = api_key
        self.temperature = temperature
        self.max_tokens = max_tokens
        super().__init__()

    def __call__(self, payload: APIPayload) -> LLMOutput:
        """Make a call to the model and return the parsed response."""
        response = self._call_api(payload)
        return self._parse_response(response)

    @abstractmethod
    def _call_api(self, payload: APIPayload) -> Any:
        """Make a call to the model API and return the raw response."""
        pass

    @abstractmethod
    def _parse_response(self, response: Any) -> LLMOutput:
        """Parse the raw response from the model API and return a structured response."""
        pass

class AgentlabAction:
    """
    Collection of utility function to convert tool calls to Agentlab action format.
    """

    @staticmethod
    def convert_toolcall_to_agentlab_action_format(toolcall: ToolCall) -> str:
        """Convert a tool call to an Agentlab environment action string.

        Args:
            toolcall: ToolCall object containing the name and arguments of the tool call.

        Returns:
            A string representing the action in Agentlab format i.e. python function call string.
        """

        tool_name, tool_args = toolcall.name, toolcall.arguments
        return tool_call_to_python_code(tool_name, tool_args)

    @staticmethod
    def convert_multiactions_to_agentlab_action_format(actions: list[str]) -> str | None:
        """Convert multiple actions list to a format that env supports.

        Args:
            actions: List of action strings to be joined.

        Returns:
            Joined actions separated by newlines, or None if empty.
        """
        return "\n".join(actions) if actions else None

class BaseModelWithPricing(TrackAPIPricingMixin, BaseResponseModel):
    pass

class OpenAIResponseModel(BaseModelWithPricing):
    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = 100,
    ):
        self.action_space_as_tools = True  # this should be a config
        super().__init__(  # This is passed to BaseModel
            model_name=model_name, api_key=api_key, temperature=temperature, max_tokens=max_tokens
        )
        client_args = {}
        if base_url is not None:
            client_args["base_url"] = base_url
        if api_key is not None:
            client_args["api_key"] = api_key
        self.client = OpenAI(**client_args)
        # Init pricing tracker after super() so that all attributes have been set.
        self.init_pricing_tracker(pricing_api="openai")  # Use the PricingMixin

    def _call_api(self, payload: APIPayload) -> "OpenAIResponseObject":

        input = []
        for msg in payload.messages:
            input.extend(msg.prepare_message())
        api_params: Dict[str, Any] = {
            "model": self.model_name,
            "input": input,
        }
        # Not all Open AI models support these parameters (example: o3), so we check if they are set.
        if self.temperature is not None:
            api_params["temperature"] = self.temperature
        if self.max_tokens is not None:
            api_params["max_output_tokens"] = self.max_tokens
        if payload.tools is not None:
            api_params["tools"] = payload.tools
        if payload.tool_choice is not None and payload.force_call_tool is None:
            api_params["tool_choice"] = (
                "required" if payload.tool_choice in ("required", "any") else payload.tool_choice
            )
        if payload.force_call_tool is not None:
            api_params["tool_choice"] = {"type": "function", "name": payload.force_call_tool}

        response = call_openai_api_with_retries(
            self.client.responses.create,
            api_params,
        )

        return response

    def _parse_response(self, response: "OpenAIResponseObject") -> LLMOutput:
        """Parse the raw response from the OpenAI Responses API."""

        think_output = self._extract_thinking_content_from_response(response)
        toolcalls = self._extract_tool_calls_from_response(response)

        if self.action_space_as_tools:
            env_action = self._extract_env_action
# ... truncated ...


ToolCall

from .llm_utils import (
    call_anthropic_api_with_retries,
    call_openai_api_with_retries,
)
from .tracking import TrackAPIPricingMixin

"""This module contains utlity classes for building input messages and interacting with LLM APIs.
It includes:
    1. Message Builder for building input messages
    2. Base Reponse class for different LLM APIs (OpenAI, Anthropic, etc.)
    3. Factory classes (inherits from BaseModelArgs) for creating instances of LLM Response models.
"""

logger = logging.getLogger(__name__)

ContentItem = Dict[str, Any]
Message = Dict[str, Union[str, List[ContentItem]]]

@dataclass
class ToolCall:
    """Represents a tool call made by the LLM.
    Attributes:
    name: Name of the tool called.
    arguments: Arguments passed to the tool.
    raw_call: The raw call object from the LLM API.
    tool_response: Output of the tool call goes here. It can be only one content item.
    """

    name: str = field(default=None)
    arguments: Dict[str, Any] = field(default_factory=dict)
    raw_call: Any = field(default=None)
    tool_response: ContentItem = None

    @property
    def is_response_set(self) -> bool:
        """Check if the tool response is set."""
        return self.tool_response is not None

    def response_text(self, text: str) -> "MessageBuilder":
        self.tool_response = {"text": text}
        return self

    def response_image(self, image: str) -> "MessageBuilder":
        self.tool_response = {"image": image}
        return self

    def __repr__(self):
        return f"ToolCall(name={self.name}, arguments={self.arguments})"

@dataclass
class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self

    @property
    def all_responses_set(self) -> bool: