main

"""
Note: This script is a convenience script to launch experiments instead of using
the command line.

Copy this script and modify at will, but don't push your changes to the
repository.
"""

import logging

from agentlab.agents.generic_agent import (
    AGENT_LLAMA3_70B,
    AGENT_LLAMA31_70B,
    RANDOM_SEARCH_AGENT,
    AGENT_4o,
    AGENT_4o_MINI,
    AGENT_o3_MINI,
    AGENT_37_SONNET,
    AGENT_CLAUDE_SONNET_35,
    AGENT_GPT5_MINI,
)
from agentlab.experiments.study import Study

logging.getLogger().setLevel(logging.INFO)

# choose your agent or provide a new agent
agent_args = [AGENT_4o_MINI]
# agent_args = [AGENT_4o]

# ## select the benchmark to run on
benchmark = "miniwob_tiny_test"
# benchmark = "miniwob"
# benchmark = "workarena_l1"
# benchmark = "workarena_l2"
# benchmark = "workarena_l3"
# benchmark = "webarena"

# Set reproducibility_mode = True for reproducibility
# this will "ask" agents to be deterministic. Also, it will prevent you from launching if you have
# local changes. For your custom agents you need to implement set_reproducibility_mode
reproducibility_mode = False

# Set relaunch = True to relaunch an existing study, this will continue incomplete
# experiments and relaunch errored experiments
relaunch = False

## Number of parallel jobs
n_jobs = 4  # Make sure to use 1 job when debugging in VSCode
# n_jobs = -1  # to use all available cores

if __name__ == "__main__":  # necessary for dask backend

    if reproducibility_mode:
        [a.set_reproducibility_mode() for a in agent_args]

    if relaunch:
        #  relaunch an existing study
        study = Study.load_most_recent(contains=None)
        study.find_incomplete(include_errors=True)

    else:
        study = Study(agent_args, benchmark, logging_level_stdout=logging.WARNING)

    study.run(
        n_jobs=n_jobs,
        parallel_backend="ray",
        strict_reproducibility=reproducibility_mode,
        n_relaunch=3,
    )

    if reproducibility_mode:
        study.append_to_journal(strict_reproducibility=True)


"""
Note: This script is a convenience script to launch experiments instead of using
the command line.

Copy this script and modify at will, but don't push your changes to the
repository.
"""

import logging
from copy import deepcopy

import bgym

from agentlab.agents.tool_use_agent.tool_use_agent import (
    DEFAULT_PROMPT_CONFIG,
    GPT_4_1,
    ToolUseAgentArgs,
)
from agentlab.experiments.study import Study

logging.getLogger().setLevel(logging.INFO)

config = deepcopy(DEFAULT_PROMPT_CONFIG)
# config.keep_last_n_obs = 1
config.obs.use_som = True

agent_configs = [
    ToolUseAgentArgs(
        model_args=GPT_4_1,
        config=config,
    ),
    # ToolUseAgentArgs(
    #     model_args=GPT_4_1,
    #     config=config,
    # ),
]

for agent_config in agent_configs:
    agent_config.config.action_subsets = ("workarena",)  # use the workarena action set

# ## select the benchmark to run on
# benchmark = "miniwob_tiny_test"
benchmark = "workarena_l1"

benchmark = bgym.DEFAULT_BENCHMARKS[benchmark](n_repeats=4)  # type: bgym.Benchmark
benchmark = benchmark.subset_from_glob("task_name", "*create*")

# for env_args in benchmark.env_args_list:
#     print(env_args.task_name)
#     env_args.max_steps = 15

relaunch = False

## Number of parallel jobs
n_jobs = 10  # Make sure to use 1 job when debugging in VSCode
parallel_backend = "ray"
# parallel_backend = "sequential"  # activate sequential backend for debugging in VSCode

if __name__ == "__main__":  # necessary for dask backend

    if relaunch:
        #  relaunch an existing study
        study = Study.load_most_recent(contains=None)
        study.find_incomplete(include_errors=True)

    else:
        study = Study(agent_configs, benchmark, logging_level_stdout=logging.WARNING)

    study.run(
        n_jobs=n_jobs,
        parallel_backend=parallel_backend,  # "ray", "joblib" or "sequential"
        strict_reproducibility=False,
        n_relaunch=3,
    )


import os
from pathlib import Path
from agentlab.experiments.study import Study

base_dir = "/home/toolkit/ui_copilot_results"

exp_paths = [
    "2025-01-31_22-08-34_genericagent-o3-mini-2025-01-31-on-workarena-l1",
    #  '2025-02-02_01-53-45_genericagent-openai-o1-mini-2024-09-12-on-workarena-l1',
    "2025-02-02_01-55-04_genericagent-openai-o1-mini-2024-09-12-on-workarena-l1",
]
full_paths = [os.path.join(base_dir, exp_path) for exp_path in exp_paths]

for full_path in full_paths:
    study = Study.load(Path(full_path))

    study.append_to_journal(strict_reproducibility=False)


import subprocess
import sys
from pathlib import Path

import pytest

@pytest.mark.pricy
def test_main_script_execution():
    # this should trigger agent_4o_mini on miniwob_tiny_test unless this was
    # reconfigured differently.
    path = Path(__file__).parent.parent / "main.py"

    sys.path.insert(0, str(path.parent))

    # just make sure it's in the right state
    main = __import__(path.stem)
    assert main.benchmark == "miniwob_tiny_test"
    assert main.reproducibility_mode == False
    assert main.relaunch == False
    assert main.n_jobs <= 10

    result = subprocess.run(["python", str(path)], capture_output=True, text=True, timeout=5 * 60)
    assert result.returncode == 0

if __name__ == "__main__":
    test_main_script_execution()


from agentlab.ui_assistant import make_exp_args
from agentlab.agents.generic_agent import AGENT_4o

def test_make_exp_args():
    """Basic unit test to detect refactoring errors."""
    exp_args = make_exp_args(AGENT_4o, "https://www.google.com")

    assert exp_args.agent_args.flags.action.demo_mode == "default"


import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import anthropic

client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

def make_request(messages):
    response = client.messages.create(
        model="claude-3-5-sonnet-20241022", max_tokens=10, messages=messages
    )
    return response.usage

def make_message(text):
    return {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": text,
            }
        ],
    }

def add_cache_control(message: dict, cache_type="ephemeral"):
    message["content"][0]["cache_control"] = {"type": cache_type}

def remove_cache_control(message: dict):
    if "cache_control" in message["content"][0]:
        del message["content"][0]["cache_control"]

def test_rate_limit_single(thread_id):
    # Create ~100k token message that will be cached
    big_text = "This is a large block of text for caching. " * 10000  # ~100k tokens
    medium_text = "This is a large block of text for caching. " * 2000  # ~10k tokens

    print(f"Thread {thread_id}: Starting rate limit test with cached content...")

    # Rebuild conversation each time (simulating web agent)
    messages = []

    # Add all previous conversation turns
    for i in range(5):
        if i == 0:
            messages.append(make_message(big_text))
            t0 = time.time()
        else:
            messages.append(make_message(medium_text))
        add_cache_control(messages[-1])
        try:
            usage = make_request(messages)
            dt = time.time() - t0
            print(f"{dt:.2f}: Thread {thread_id}: {usage}")
        except Exception as e:
            print(f"Thread {thread_id}: Error - {e}")
            break
        remove_cache_control(messages[-1])

def test_rate_limit_parallel(num_threads=3):
    print(f"Starting parallel rate limit test with {num_threads} threads...")

    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        futures = [executor.submit(test_rate_limit_single, i) for i in range(num_threads)]

        for future in as_completed(futures):
            try:
                future.result()
            except Exception as e:
                print(f"Thread completed with error: {e}")

def test_rate_limit():
    # Original single-threaded version
    test_rate_limit_single(0)

if __name__ == "__main__":
    # Use parallel version to quickly exhaust rate limits
    test_rate_limit_parallel(num_threads=3)

    # Or use original single-threaded version
    # test_rate_limit()


import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import anthropic

client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

def make_request(messages):
    response = client.messages.create(
        model="claude-3-5-sonnet-20241022", max_tokens=10, messages=messages
    )
    return response.usage

def make_message(text):
    return {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": text,
            }
        ],
    }

def add_cache_control(message: dict, cache_type="ephemeral"):
    message["content"][0]["cache_control"] = {"type": cache_type}

def remove_cache_control(message: dict):
    if "cache_control" in message["content"][0]:


import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import anthropic

client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

def make_request(messages):
    response = client.messages.create(
        model="claude-3-5-sonnet-20241022", max_tokens=10, messages=messages
    )
    return response.usage

def make_message(text):
    return {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": text,
            }
        ],
    }

def add_cache_control(message: dict, cache_type="ephemeral"):
    message["content"][0]["cache_control"] = {"type": cache_type}

def remove_cache_control(message: dict):
    if "cache_control" in message["content"][0]:
        del message["content"][0]["cache_control"]

def test_rate_limit_single(thread_id):
    # Create ~100k token message that will be cached
    big_text = "This is a large block of text for caching. " * 10000  # ~100k tokens
    medium_text = "This is a large block of text for caching. " * 2000  # ~10k tokens

    print(f"Thread {thread_id}: Starting rate limit test with cached content...")

    # Rebuild conversation each time (simulating web agent)
    messages = []


def make_request(messages):
    response = client.messages.create(
        model="claude-3-5-sonnet-20241022", max_tokens=10, messages=messages
    )
    return response.usage

def make_message(text):
    return {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": text,
            }
        ],
    }

def add_cache_control(message: dict, cache_type="ephemeral"):
    message["content"][0]["cache_control"] = {"type": cache_type}

def remove_cache_control(message: dict):
    if "cache_control" in message["content"][0]:
        del message["content"][0]["cache_control"]

def test_rate_limit_single(thread_id):
    # Create ~100k token message that will be cached
    big_text = "This is a large block of text for caching. " * 10000  # ~100k tokens
    medium_text = "This is a large block of text for caching. " * 2000  # ~10k tokens

    print(f"Thread {thread_id}: Starting rate limit test with cached content...")

    # Rebuild conversation each time (simulating web agent)
    messages = []

    # Add all previous conversation turns
    for i in range(5):
        if i == 0:


    )
    return response.usage

def make_message(text):
    return {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": text,
            }
        ],
    }

def add_cache_control(message: dict, cache_type="ephemeral"):
    message["content"][0]["cache_control"] = {"type": cache_type}

def remove_cache_control(message: dict):
    if "cache_control" in message["content"][0]:
        del message["content"][0]["cache_control"]

def test_rate_limit_single(thread_id):
    # Create ~100k token message that will be cached
    big_text = "This is a large block of text for caching. " * 10000  # ~100k tokens
    medium_text = "This is a large block of text for caching. " * 2000  # ~10k tokens

    print(f"Thread {thread_id}: Starting rate limit test with cached content...")

    # Rebuild conversation each time (simulating web agent)
    messages = []

    # Add all previous conversation turns
    for i in range(5):
        if i == 0:
            messages.append(make_message(big_text))
            t0 = time.time()
        else:
            messages.append(make_message(medium_text))
        add_cache_control(messages[-1])


    return {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": text,
            }
        ],
    }

def add_cache_control(message: dict, cache_type="ephemeral"):
    message["content"][0]["cache_control"] = {"type": cache_type}

def remove_cache_control(message: dict):
    if "cache_control" in message["content"][0]:
        del message["content"][0]["cache_control"]

def test_rate_limit_single(thread_id):
    # Create ~100k token message that will be cached
    big_text = "This is a large block of text for caching. " * 10000  # ~100k tokens
    medium_text = "This is a large block of text for caching. " * 2000  # ~10k tokens

    print(f"Thread {thread_id}: Starting rate limit test with cached content...")

    # Rebuild conversation each time (simulating web agent)
    messages = []

    # Add all previous conversation turns
    for i in range(5):
        if i == 0:
            messages.append(make_message(big_text))
            t0 = time.time()
        else:
            messages.append(make_message(medium_text))
        add_cache_control(messages[-1])
        try:
            usage = make_request(messages)
            dt = time.time() - t0
            print(f"{dt:.2f}: Thread {thread_id}: {usage}")
        except Exception as e:
            print(f"Thread {thread_id}: Error - {e}")
            break
        remove_cache_control(messages[-1])

def test_rate_limit_parallel(num_threads=3):
    print(f"Starting parallel rate limit test with {num_threads} threads...")

    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        futures = [executor.submit(test_rate_limit_single, i) for i in range(num_threads)]

        for future in as_completed(futures):
            try:
                future.result()
            except Exception as e:
                print(f"Thread completed with error: {e}")

def test_rate_limit():
    # Original single-threaded version
    test_rate_limit_single(0)


    messages = []

    # Add all previous conversation turns
    for i in range(5):
        if i == 0:
            messages.append(make_message(big_text))
            t0 = time.time()
        else:
            messages.append(make_message(medium_text))
        add_cache_control(messages[-1])
        try:
            usage = make_request(messages)
            dt = time.time() - t0
            print(f"{dt:.2f}: Thread {thread_id}: {usage}")
        except Exception as e:
            print(f"Thread {thread_id}: Error - {e}")
            break
        remove_cache_control(messages[-1])

def test_rate_limit_parallel(num_threads=3):
    print(f"Starting parallel rate limit test with {num_threads} threads...")

    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        futures = [executor.submit(test_rate_limit_single, i) for i in range(num_threads)]

        for future in as_completed(futures):
            try:
                future.result()
            except Exception as e:
                print(f"Thread completed with error: {e}")

def test_rate_limit():
    # Original single-threaded version
    test_rate_limit_single(0)

if __name__ == "__main__":
    # Use parallel version to quickly exhaust rate limits
    test_rate_limit_parallel(num_threads=3)

    # Or use original single-threaded version
    # test_rate_limit()


            print(f"{dt:.2f}: Thread {thread_id}: {usage}")
        except Exception as e:
            print(f"Thread {thread_id}: Error - {e}")
            break
        remove_cache_control(messages[-1])

def test_rate_limit_parallel(num_threads=3):
    print(f"Starting parallel rate limit test with {num_threads} threads...")

    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        futures = [executor.submit(test_rate_limit_single, i) for i in range(num_threads)]

        for future in as_completed(futures):
            try:
                future.result()
            except Exception as e:
                print(f"Thread completed with error: {e}")

def test_rate_limit():
    # Original single-threaded version
    test_rate_limit_single(0)

if __name__ == "__main__":
    # Use parallel version to quickly exhaust rate limits
    test_rate_limit_parallel(num_threads=3)

    # Or use original single-threaded version
    # test_rate_limit()


from PIL import Image

from agentlab.analyze import overlay_utils

def test_parse_function_calls():

    test_code = """
mouse_click(34, 59)
fill("a234", "test")
click('b123', button="right", modifiers=["Shift", "Control"])
select_option("c456", ["option1", "option2"])
"""

    result = overlay_utils.parse_function_calls(test_code)

    assert result[1].function_name == "mouse_click"
    assert result[1].name == "y"
    assert test_code[result[1].start_index : result[1].stop_index] == "59"

    assert result[8].function_name == "select_option"
    assert result[8].name == "options"
    assert test_code[result[8].start_index : result[8].stop_index] == '["option1", "option2"]'

def test_filtering_args():
    test_code = """
mouse_click(34, 59)
fill("a234", "test")
mouse_drag_and_drop(34, 59, to_x=100, to_y=200)
drag_and_drop("a123", "b456")
"""
    result = overlay_utils.parse_function_calls(test_code)
    args = overlay_utils.find_bids_and_xy_pairs(result)

    assert len(args) == 6  # Expecting 4 args: 2 mouse clicks, 1 fill, 1 select_option

    assert args[0].function_name == "mouse_click"
    assert args[0].name == "xy"
    assert args[0].value == (34.0, 59.0)
    assert test_code[args[0].start_index : args[0].stop_index] == "34, 59"

    assert args[2].name == "from_xy"
    assert args[3].name == "to_xy"
    assert test_code[args[3].start_index : args[3].stop_index] == "to_x=100, to_y=200"

def manual_eval():
    """Manual test function that displays the resulting image."""
    import matplotlib.pyplot as plt

    # Create a white test image
    img = Image.new("RGB", (400, 300), "white")

    # Test action string with multiple function calls
    action_string = """mouse_click(100, 150)
fill("search_box", "hello world")
click("submit_btn")"""

    # Mock properties mapping bids to bounding boxes
    properties = {
        "search_box": {"bbox": (50, 50, 100, 50)},
        "submit_btn": {"bbox": (150, 100, 120, 30)},
    }

    # Annotate the image and get colored HTML
    html_result = overlay_utils.annotate_action(img, action_string, properties, colormap="tab10")

    # Display result
    plt.figure(figsize=(10, 6))
    plt.imshow(img)
    plt.axis("off")
    plt.show()

    print("HTML with colored arguments:")
    print(html_result)
    print("\nManual test completed!")

if __name__ == "__main__":
    manual_eval()


from PIL import Image

from agentlab.analyze import overlay_utils

def test_parse_function_calls():

    test_code = """
mouse_click(34, 59)
fill("a234", "test")
click('b123', button="right", modifiers=["Shift", "Control"])
select_option("c456", ["option1", "option2"])
"""

    result = overlay_utils.parse_function_calls(test_code)

    assert result[1].function_name == "mouse_click"
    assert result[1].name == "y"
    assert test_code[result[1].start_index : result[1].stop_index] == "59"

    assert result[8].function_name == "select_option"
    assert result[8].name == "options"
    assert test_code[result[8].start_index : result[8].stop_index] == '["option1", "option2"]'

def test_filtering_args():
    test_code = """
mouse_click(34, 59)
fill("a234", "test")
mouse_drag_and_drop(34, 59, to_x=100, to_y=200)
drag_and_drop("a123", "b456")
"""
    result = overlay_utils.parse_function_calls(test_code)
    args = overlay_utils.find_bids_and_xy_pairs(result)

    assert len(args) == 6  # Expecting 4 args: 2 mouse clicks, 1 fill, 1 select_option

    assert args[0].function_name == "mouse_click"
    assert args[0].name == "xy"
    assert args[0].value == (34.0, 59.0)
    assert test_code[args[0].start_index : args[0].stop_index] == "34, 59"

    assert args[2].name == "from_xy"


def test_parse_function_calls():

    test_code = """
mouse_click(34, 59)
fill("a234", "test")
click('b123', button="right", modifiers=["Shift", "Control"])
select_option("c456", ["option1", "option2"])
"""

    result = overlay_utils.parse_function_calls(test_code)

    assert result[1].function_name == "mouse_click"
    assert result[1].name == "y"
    assert test_code[result[1].start_index : result[1].stop_index] == "59"

    assert result[8].function_name == "select_option"
    assert result[8].name == "options"
    assert test_code[result[8].start_index : result[8].stop_index] == '["option1", "option2"]'

def test_filtering_args():
    test_code = """
mouse_click(34, 59)
fill("a234", "test")
mouse_drag_and_drop(34, 59, to_x=100, to_y=200)
drag_and_drop("a123", "b456")
"""
    result = overlay_utils.parse_function_calls(test_code)
    args = overlay_utils.find_bids_and_xy_pairs(result)

    assert len(args) == 6  # Expecting 4 args: 2 mouse clicks, 1 fill, 1 select_option

    assert args[0].function_name == "mouse_click"
    assert args[0].name == "xy"
    assert args[0].value == (34.0, 59.0)
    assert test_code[args[0].start_index : args[0].stop_index] == "34, 59"

    assert args[2].name == "from_xy"
    assert args[3].name == "to_xy"
    assert test_code[args[3].start_index : args[3].stop_index] == "to_x=100, to_y=200"

def manual_eval():
    """Manual test function that displays the resulting image."""
    import matplotlib.pyplot as plt

    # Create a white test image
    img = Image.new("RGB", (400, 300), "white")

    # Test action string with multiple function calls
    action_string = """mouse_click(100, 150)
fill("search_box", "hello world")
click("submit_btn")"""

    # Mock properties mapping bids to bounding boxes
    properties = {
        "search_box": {"bbox": (50, 50, 100, 50)},
        "submit_btn": {"bbox": (150, 100, 120, 30)},
    }


mouse_click(34, 59)
fill("a234", "test")
mouse_drag_and_drop(34, 59, to_x=100, to_y=200)
drag_and_drop("a123", "b456")
"""
    result = overlay_utils.parse_function_calls(test_code)
    args = overlay_utils.find_bids_and_xy_pairs(result)

    assert len(args) == 6  # Expecting 4 args: 2 mouse clicks, 1 fill, 1 select_option

    assert args[0].function_name == "mouse_click"
    assert args[0].name == "xy"
    assert args[0].value == (34.0, 59.0)
    assert test_code[args[0].start_index : args[0].stop_index] == "34, 59"

    assert args[2].name == "from_xy"
    assert args[3].name == "to_xy"
    assert test_code[args[3].start_index : args[3].stop_index] == "to_x=100, to_y=200"

def manual_eval():
    """Manual test function that displays the resulting image."""
    import matplotlib.pyplot as plt

    # Create a white test image
    img = Image.new("RGB", (400, 300), "white")

    # Test action string with multiple function calls
    action_string = """mouse_click(100, 150)
fill("search_box", "hello world")
click("submit_btn")"""

    # Mock properties mapping bids to bounding boxes
    properties = {
        "search_box": {"bbox": (50, 50, 100, 50)},
        "submit_btn": {"bbox": (150, 100, 120, 30)},
    }

    # Annotate the image and get colored HTML
    html_result = overlay_utils.annotate_action(img, action_string, properties, colormap="tab10")

    # Display result
    plt.figure(figsize=(10, 6))
    plt.imshow(img)
    plt.axis("off")
    plt.show()

    print("HTML with colored arguments:")
    print(html_result)
    print("\nManual test completed!")

if __name__ == "__main__":
    manual_eval()


from pathlib import Path
import shutil
import tempfile

import pandas as pd
from agentlab.analyze.inspect_results import get_study_summary

def test_get_study_summary():

    with tempfile.TemporaryDirectory() as tmp_dir:
        study_dir = Path(tmp_dir) / "test_study"

        study_dir_original = Path(__file__).parent.parent / "data" / "test_study"

        # recursively copy the study to the temp dir using shutil
        shutil.copytree(study_dir_original, study_dir)

        sentinel = {}

        summary = get_study_summary(study_dir, sentinel=sentinel)
        assert isinstance(summary, pd.DataFrame)
        assert sentinel["from_cache"] == False

        summary = get_study_summary(study_dir, sentinel=sentinel)
        assert isinstance(summary, pd.DataFrame)
        assert sentinel["from_cache"] == True

        summary = get_study_summary(study_dir, ignore_cache=True, sentinel=sentinel)
        assert isinstance(summary, pd.DataFrame)
        assert sentinel["from_cache"] == False

if __name__ == "__main__":
    test_get_study_summary()


import pytest

from agentlab.llm.chat_api import HuggingFaceURLChatModel, make_system_message, make_user_message
from agentlab.llm.llm_utils import download_and_save_model
from agentlab.llm.prompt_templates import STARCHAT_PROMPT_TEMPLATE

# TODO(optimass): figure out a good model for all tests

@pytest.mark.skip(reason="Requires a local model checkpoint")
def test_CustomLLMChatbot_locally():
    # model_path = "google/flan-t5-base"  # remote model on HuggingFace Hub
    model_path = "/mnt/ui_copilot/data_rw/models/starcoderbase-1b-ft"  # local model in shared volum

    chatbot = HuggingFaceURLChatModel(model_path=model_path, temperature=1e-3)

    messages = [
        make_system_message("Please tell me back the following word: "),
        make_user_message("bird"),
    ]

    answer = chatbot(messages)

    print(answer.content)

@pytest.mark.skip(reason="Requires downloading a large file on disk local model checkpoint")
def test_download_and_save_model():
    model_path = "meta-llama/Llama-2-70b-chat"
    save_dir = "test_models"

    download_and_save_model(model_path, save_dir)


@pytest.mark.skip(reason="Requires a local model checkpoint")
def test_CustomLLMChatbot_locally():
    # model_path = "google/flan-t5-base"  # remote model on HuggingFace Hub
    model_path = "/mnt/ui_copilot/data_rw/models/starcoderbase-1b-ft"  # local model in shared volum

    chatbot = HuggingFaceURLChatModel(model_path=model_path, temperature=1e-3)

    messages = [
        make_system_message("Please tell me back the following word: "),
        make_user_message("bird"),
    ]

    answer = chatbot(messages)

    print(answer.content)

@pytest.mark.skip(reason="Requires downloading a large file on disk local model checkpoint")
def test_download_and_save_model():
    model_path = "meta-llama/Llama-2-70b-chat"
    save_dir = "test_models"

    download_and_save_model(model_path, save_dir)


from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
from agentlab.llm.chat_api import BaseModelArgs

def test_llm_configs():

    for _, args in CHAT_MODEL_ARGS_DICT.items():
        assert isinstance(args, BaseModelArgs)


import os
import time
from functools import partial

import pytest

import agentlab.llm.tracking as tracking
from agentlab.llm.chat_api import (
    AzureChatModel,
    OpenAIChatModel,
    OpenRouterChatModel,
    make_system_message,
    make_user_message,
)

def test_get_action_decorator():
    action, agent_info = tracking.cost_tracker_decorator(lambda x, y: call_llm())(None, None)
    assert action == "action"
    assert agent_info["stats"] == {
        "input_tokens": 1,
        "output_tokens": 1,
        "cost": 1.0,
    }

OPENROUTER_API_KEY_AVAILABLE = os.environ.get("OPENROUTER_API_KEY") is not None

OPENROUTER_MODELS = (
    "anthropic/claude-3.5-sonnet",
    "meta-llama/llama-3.1-405b-instruct",
    "meta-llama/llama-3.1-70b-instruct",
    "meta-llama/llama-3.1-8b-instruct",
    "google/gemini-pro-1.5",
)

@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available")
def test_get_pricing_openrouter():
    pricing = tracking.get_pricing_openrouter()
    assert isinstance(pricing, dict)
    assert all(isinstance(v, dict) for v in pricing.values())
    for model in OPENROUTER_MODELS:
        assert model in pricing
        assert isinstance(pricing[model], dict)
        assert all(isinstance(v, float) for v in pricing[model].values())

def test_get_pricing_openai():
    pricing = tracking.get_pricing_openai()
    assert isinstance(pricing, dict)
    assert all("prompt" in pricing[model] and "completion" in pricing[model] for model in pricing)
    assert all(isinstance(pricing[model]["prompt"], float) for model in pricing)
    assert all(isinstance(pricing[model]["completion"], float) for model in pricing)

def call_llm():
    if hasattr(tracking.TRACKER, "instance") and isinstance(
        tracking.TRACKER.instance, tracking.LLMTracker
    ):
        tracking.TRACKER.instance(1, 1, 1)
    return "action", {"stats": {}}

def test_tracker():
    with tracking.set_tracker() as tracker:
        _, _ = call_llm()

    assert tracker.stats["cost"] == 1

def test_imbricate_trackers():
    with tracking.set_tracker() as tracker4:
        with tracking.set_tracker() as tracker1:
            _, _ = call_llm()
        with tracking.set_tracker() as tracker3:
            _, _ = call_llm()
            _, _ = call_llm()
            with tracking.set_tracker() as tracker1bis:
                _, _ = call_llm()

    assert tracker1.stats["cost"] == 1
    assert tracker1bis.stats["cost"] == 1
    assert tracker3.stats["cost"] == 3
    assert tracker4.stats["cost"] == 4

def test_threaded_trackers():
    """thread_2 occurs in the middle of thread_1, results should be separate."""
    import threading

    def thread_1(results=None):
        with tracking.set_tracker() as tracker:
            time.sleep(1)
            _, _ = call_llm()
            time.sleep(1)
        results[0] = tracker.stats

    def thread_2(results=None):
        time.sleep(1)
        with tracking.set_tracker() as tracker:
            _, _ = call_llm()
        results[1] = tracker.stats

    results = [None] * 2
    threads = [
        threading.Thread(target=partial(thread_1, results=results)),
        threading.Thread(target=partial(thread_2, results=results)),
    ]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()

    assert all(result["cost"] == 1 for result in results)

OPENAI_API_KEY_AVAILABLE = os.environ.get("OPENAI_API_KEY") is not None

@pytest.mark.pricy
@pytest.mark.skipif(not OPENAI_API_KEY_AVAILABLE, reason="OpenAI API key is not available")
def test_openai_chat_model():
    chat_model = OpenAIChatModel("gpt-4o-mini")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0

AZURE_OPENAI_API_KEY_AVAILABLE = (
    os.environ.get("AZURE_OPENAI_API_KEY") is not None
    and os.environ.get("AZURE_OPENAI_ENDPOINT") is not None
)

@pytest.mark.pricy
@pytest.mark.skipif(
    not AZURE_OPENAI_API_KEY_AVAILABLE, reason="Azure OpenAI API key is not available"
)
def test_azure_chat_model():
    chat_model = AzureChatModel(model_name="gpt-4.1-nano", deployment_name="gpt-4.1-nano")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available")
def test_openrouter_chat_model():
    chat_model = OpenRouterChatModel("openai/gpt-4o-mini")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0


import os
import time
from functools import partial

import pytest

import agentlab.llm.tracking as tracking
from agentlab.llm.chat_api import (
    AzureChatModel,
    OpenAIChatModel,
    OpenRouterChatModel,
    make_system_message,
    make_user_message,
)

def test_get_action_decorator():
    action, agent_info = tracking.cost_tracker_decorator(lambda x, y: call_llm())(None, None)
    assert action == "action"
    assert agent_info["stats"] == {
        "input_tokens": 1,
        "output_tokens": 1,
        "cost": 1.0,
    }

OPENROUTER_API_KEY_AVAILABLE = os.environ.get("OPENROUTER_API_KEY") is not None

OPENROUTER_MODELS = (
    "anthropic/claude-3.5-sonnet",
    "meta-llama/llama-3.1-405b-instruct",
    "meta-llama/llama-3.1-70b-instruct",
    "meta-llama/llama-3.1-8b-instruct",
    "google/gemini-pro-1.5",
)

@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available")
def test_get_pricing_openrouter():
    pricing = tracking.get_pricing_openrouter()
    assert isinstance(pricing, dict)
    assert all(isinstance(v, dict) for v in pricing.values())
    for model in OPENROUTER_MODELS:
        assert model in pricing


    assert action == "action"
    assert agent_info["stats"] == {
        "input_tokens": 1,
        "output_tokens": 1,
        "cost": 1.0,
    }

OPENROUTER_API_KEY_AVAILABLE = os.environ.get("OPENROUTER_API_KEY") is not None

OPENROUTER_MODELS = (
    "anthropic/claude-3.5-sonnet",
    "meta-llama/llama-3.1-405b-instruct",
    "meta-llama/llama-3.1-70b-instruct",
    "meta-llama/llama-3.1-8b-instruct",
    "google/gemini-pro-1.5",
)

@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available")
def test_get_pricing_openrouter():
    pricing = tracking.get_pricing_openrouter()
    assert isinstance(pricing, dict)
    assert all(isinstance(v, dict) for v in pricing.values())
    for model in OPENROUTER_MODELS:
        assert model in pricing
        assert isinstance(pricing[model], dict)
        assert all(isinstance(v, float) for v in pricing[model].values())

def test_get_pricing_openai():
    pricing = tracking.get_pricing_openai()
    assert isinstance(pricing, dict)
    assert all("prompt" in pricing[model] and "completion" in pricing[model] for model in pricing)
    assert all(isinstance(pricing[model]["prompt"], float) for model in pricing)
    assert all(isinstance(pricing[model]["completion"], float) for model in pricing)

def call_llm():
    if hasattr(tracking.TRACKER, "instance") and isinstance(
        tracking.TRACKER.instance, tracking.LLMTracker
    ):
        tracking.TRACKER.instance(1, 1, 1)
    return "action", {"stats": {}}

def test_tracker():
    with tracking.set_tracker() as tracker:


OPENROUTER_MODELS = (
    "anthropic/claude-3.5-sonnet",
    "meta-llama/llama-3.1-405b-instruct",
    "meta-llama/llama-3.1-70b-instruct",
    "meta-llama/llama-3.1-8b-instruct",
    "google/gemini-pro-1.5",
)

@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available")
def test_get_pricing_openrouter():
    pricing = tracking.get_pricing_openrouter()
    assert isinstance(pricing, dict)
    assert all(isinstance(v, dict) for v in pricing.values())
    for model in OPENROUTER_MODELS:
        assert model in pricing
        assert isinstance(pricing[model], dict)
        assert all(isinstance(v, float) for v in pricing[model].values())

def test_get_pricing_openai():
    pricing = tracking.get_pricing_openai()
    assert isinstance(pricing, dict)
    assert all("prompt" in pricing[model] and "completion" in pricing[model] for model in pricing)
    assert all(isinstance(pricing[model]["prompt"], float) for model in pricing)
    assert all(isinstance(pricing[model]["completion"], float) for model in pricing)

def call_llm():
    if hasattr(tracking.TRACKER, "instance") and isinstance(
        tracking.TRACKER.instance, tracking.LLMTracker
    ):
        tracking.TRACKER.instance(1, 1, 1)
    return "action", {"stats": {}}

def test_tracker():
    with tracking.set_tracker() as tracker:
        _, _ = call_llm()

    assert tracker.stats["cost"] == 1

def test_imbricate_trackers():
    with tracking.set_tracker() as tracker4:
        with tracking.set_tracker() as tracker1:


@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available")
def test_get_pricing_openrouter():
    pricing = tracking.get_pricing_openrouter()
    assert isinstance(pricing, dict)
    assert all(isinstance(v, dict) for v in pricing.values())
    for model in OPENROUTER_MODELS:
        assert model in pricing
        assert isinstance(pricing[model], dict)
        assert all(isinstance(v, float) for v in pricing[model].values())

def test_get_pricing_openai():
    pricing = tracking.get_pricing_openai()
    assert isinstance(pricing, dict)
    assert all("prompt" in pricing[model] and "completion" in pricing[model] for model in pricing)
    assert all(isinstance(pricing[model]["prompt"], float) for model in pricing)
    assert all(isinstance(pricing[model]["completion"], float) for model in pricing)

def call_llm():
    if hasattr(tracking.TRACKER, "instance") and isinstance(
        tracking.TRACKER.instance, tracking.LLMTracker
    ):
        tracking.TRACKER.instance(1, 1, 1)
    return "action", {"stats": {}}

def test_tracker():
    with tracking.set_tracker() as tracker:
        _, _ = call_llm()

    assert tracker.stats["cost"] == 1

def test_imbricate_trackers():
    with tracking.set_tracker() as tracker4:
        with tracking.set_tracker() as tracker1:
            _, _ = call_llm()
        with tracking.set_tracker() as tracker3:
            _, _ = call_llm()
            _, _ = call_llm()
            with tracking.set_tracker() as tracker1bis:
                _, _ = call_llm()

    assert tracker1.stats["cost"] == 1


        assert isinstance(pricing[model], dict)
        assert all(isinstance(v, float) for v in pricing[model].values())

def test_get_pricing_openai():
    pricing = tracking.get_pricing_openai()
    assert isinstance(pricing, dict)
    assert all("prompt" in pricing[model] and "completion" in pricing[model] for model in pricing)
    assert all(isinstance(pricing[model]["prompt"], float) for model in pricing)
    assert all(isinstance(pricing[model]["completion"], float) for model in pricing)

def call_llm():
    if hasattr(tracking.TRACKER, "instance") and isinstance(
        tracking.TRACKER.instance, tracking.LLMTracker
    ):
        tracking.TRACKER.instance(1, 1, 1)
    return "action", {"stats": {}}

def test_tracker():
    with tracking.set_tracker() as tracker:
        _, _ = call_llm()

    assert tracker.stats["cost"] == 1

def test_imbricate_trackers():
    with tracking.set_tracker() as tracker4:
        with tracking.set_tracker() as tracker1:
            _, _ = call_llm()
        with tracking.set_tracker() as tracker3:
            _, _ = call_llm()
            _, _ = call_llm()
            with tracking.set_tracker() as tracker1bis:
                _, _ = call_llm()

    assert tracker1.stats["cost"] == 1
    assert tracker1bis.stats["cost"] == 1
    assert tracker3.stats["cost"] == 3
    assert tracker4.stats["cost"] == 4

def test_threaded_trackers():
    """thread_2 occurs in the middle of thread_1, results should be separate."""


    assert all("prompt" in pricing[model] and "completion" in pricing[model] for model in pricing)
    assert all(isinstance(pricing[model]["prompt"], float) for model in pricing)
    assert all(isinstance(pricing[model]["completion"], float) for model in pricing)

def call_llm():
    if hasattr(tracking.TRACKER, "instance") and isinstance(
        tracking.TRACKER.instance, tracking.LLMTracker
    ):
        tracking.TRACKER.instance(1, 1, 1)
    return "action", {"stats": {}}

def test_tracker():
    with tracking.set_tracker() as tracker:
        _, _ = call_llm()

    assert tracker.stats["cost"] == 1

def test_imbricate_trackers():
    with tracking.set_tracker() as tracker4:
        with tracking.set_tracker() as tracker1:
            _, _ = call_llm()
        with tracking.set_tracker() as tracker3:
            _, _ = call_llm()
            _, _ = call_llm()
            with tracking.set_tracker() as tracker1bis:
                _, _ = call_llm()

    assert tracker1.stats["cost"] == 1
    assert tracker1bis.stats["cost"] == 1
    assert tracker3.stats["cost"] == 3
    assert tracker4.stats["cost"] == 4

def test_threaded_trackers():
    """thread_2 occurs in the middle of thread_1, results should be separate."""
    import threading

    def thread_1(results=None):
        with tracking.set_tracker() as tracker:
            time.sleep(1)
            _, _ = call_llm()
            time.sleep(1)
        results[0] = tracker.stats

    def thread_2(results=None):
        time.sleep(1)
        with tracking.set_tracker() as tracker:
            _, _ = call_llm()
        results[1] = tracker.stats

    results = [None] * 2


    assert tracker.stats["cost"] == 1

def test_imbricate_trackers():
    with tracking.set_tracker() as tracker4:
        with tracking.set_tracker() as tracker1:
            _, _ = call_llm()
        with tracking.set_tracker() as tracker3:
            _, _ = call_llm()
            _, _ = call_llm()
            with tracking.set_tracker() as tracker1bis:
                _, _ = call_llm()

    assert tracker1.stats["cost"] == 1
    assert tracker1bis.stats["cost"] == 1
    assert tracker3.stats["cost"] == 3
    assert tracker4.stats["cost"] == 4

def test_threaded_trackers():
    """thread_2 occurs in the middle of thread_1, results should be separate."""
    import threading

    def thread_1(results=None):
        with tracking.set_tracker() as tracker:
            time.sleep(1)
            _, _ = call_llm()
            time.sleep(1)
        results[0] = tracker.stats

    def thread_2(results=None):
        time.sleep(1)
        with tracking.set_tracker() as tracker:
            _, _ = call_llm()
        results[1] = tracker.stats

    results = [None] * 2
    threads = [
        threading.Thread(target=partial(thread_1, results=results)),
        threading.Thread(target=partial(thread_2, results=results)),
    ]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()

    assert all(result["cost"] == 1 for result in results)

OPENAI_API_KEY_AVAILABLE = os.environ.get("OPENAI_API_KEY") is not None

@pytest.mark.pricy
@pytest.mark.skipif(not OPENAI_API_KEY_AVAILABLE, reason="OpenAI API key is not available")
def test_openai_chat_model():
    chat_model = OpenAIChatModel("gpt-4o-mini")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0


        results[1] = tracker.stats

    results = [None] * 2
    threads = [
        threading.Thread(target=partial(thread_1, results=results)),
        threading.Thread(target=partial(thread_2, results=results)),
    ]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()

    assert all(result["cost"] == 1 for result in results)

OPENAI_API_KEY_AVAILABLE = os.environ.get("OPENAI_API_KEY") is not None

@pytest.mark.pricy
@pytest.mark.skipif(not OPENAI_API_KEY_AVAILABLE, reason="OpenAI API key is not available")
def test_openai_chat_model():
    chat_model = OpenAIChatModel("gpt-4o-mini")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0

AZURE_OPENAI_API_KEY_AVAILABLE = (
    os.environ.get("AZURE_OPENAI_API_KEY") is not None
    and os.environ.get("AZURE_OPENAI_ENDPOINT") is not None
)

@pytest.mark.pricy
@pytest.mark.skipif(
    not AZURE_OPENAI_API_KEY_AVAILABLE, reason="Azure OpenAI API key is not available"
)
def test_azure_chat_model():
    chat_model = AzureChatModel(model_name="gpt-4.1-nano", deployment_name="gpt-4.1-nano")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),


    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0

AZURE_OPENAI_API_KEY_AVAILABLE = (
    os.environ.get("AZURE_OPENAI_API_KEY") is not None
    and os.environ.get("AZURE_OPENAI_ENDPOINT") is not None
)

@pytest.mark.pricy
@pytest.mark.skipif(
    not AZURE_OPENAI_API_KEY_AVAILABLE, reason="Azure OpenAI API key is not available"
)
def test_azure_chat_model():
    chat_model = AzureChatModel(model_name="gpt-4.1-nano", deployment_name="gpt-4.1-nano")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available")
def test_openrouter_chat_model():
    chat_model = OpenRouterChatModel("openai/gpt-4o-mini")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0


@pytest.mark.skipif(
    not AZURE_OPENAI_API_KEY_AVAILABLE, reason="Azure OpenAI API key is not available"
)
def test_azure_chat_model():
    chat_model = AzureChatModel(model_name="gpt-4.1-nano", deployment_name="gpt-4.1-nano")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not OPENROUTER_API_KEY_AVAILABLE, reason="OpenRouter API key is not available")
def test_openrouter_chat_model():
    chat_model = OpenRouterChatModel("openai/gpt-4o-mini")
    assert chat_model.input_cost > 0
    assert chat_model.output_cost > 0

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    with tracking.set_tracker() as tracker:
        answer = chat_model(messages)
    assert "5" in answer.get("content")
    assert tracker.stats["cost"] > 0


def test_imbricate_trackers():
    with tracking.set_tracker() as tracker4:
        with tracking.set_tracker() as tracker1:
            _, _ = call_llm()
        with tracking.set_tracker() as tracker3:
            _, _ = call_llm()
            _, _ = call_llm()
            with tracking.set_tracker() as tracker1bis:
                _, _ = call_llm()

    assert tracker1.stats["cost"] == 1
    assert tracker1bis.stats["cost"] == 1
    assert tracker3.stats["cost"] == 3
    assert tracker4.stats["cost"] == 4

def test_threaded_trackers():
    """thread_2 occurs in the middle of thread_1, results should be separate."""
    import threading

    def thread_1(results=None):
        with tracking.set_tracker() as tracker:
            time.sleep(1)
            _, _ = call_llm()
            time.sleep(1)
        results[0] = tracker.stats

    def thread_2(results=None):
        time.sleep(1)
        with tracking.set_tracker() as tracker:
            _, _ = call_llm()
        results[1] = tracker.stats

    results = [None] * 2
    threads = [
        threading.Thread(target=partial(thread_1, results=results)),
        threading.Thread(target=partial(thread_2, results=results)),
    ]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()

    assert all(result["cost"] == 1 for result in results)


            with tracking.set_tracker() as tracker1bis:
                _, _ = call_llm()

    assert tracker1.stats["cost"] == 1
    assert tracker1bis.stats["cost"] == 1
    assert tracker3.stats["cost"] == 3
    assert tracker4.stats["cost"] == 4

def test_threaded_trackers():
    """thread_2 occurs in the middle of thread_1, results should be separate."""
    import threading

    def thread_1(results=None):
        with tracking.set_tracker() as tracker:
            time.sleep(1)
            _, _ = call_llm()
            time.sleep(1)
        results[0] = tracker.stats

    def thread_2(results=None):
        time.sleep(1)
        with tracking.set_tracker() as tracker:
            _, _ = call_llm()
        results[1] = tracker.stats

    results = [None] * 2
    threads = [
        threading.Thread(target=partial(thread_1, results=results)),
        threading.Thread(target=partial(thread_2, results=results)),
    ]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()

    assert all(result["cost"] == 1 for result in results)

OPENAI_API_KEY_AVAILABLE = os.environ.get("OPENAI_API_KEY") is not None

@pytest.mark.pricy
@pytest.mark.skipif(not OPENAI_API_KEY_AVAILABLE, reason="OpenAI API key is not available")
def test_openai_chat_model():


import os
from functools import partial

import pytest
from agentlab.llm.litellm_api import LiteLLMModelArgs
from agentlab.llm.response_api import APIPayload, LLMOutput

chat_api_tools = [
    {
        "type": "function",
        "name": "get_weather",
        "description": "Get the current weather in a given location.",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The location to get the weather for.",
                },
                "unit": {
                    "type": "string",
                    "enum": ["celsius", "fahrenheit"],
                    "description": "The unit of temperature.",
                },
            },
            "required": ["location"],
        },
    },
    {
        "type": "function",
        "name": "get_time",
        "description": "Get the current time in a given location.",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The location to get the time for.",
                }
            },
            "required": ["location"],
        },
    },
]

# test_config (setting name, BaseModelArgs, model_name, tools)
tool_test_configs = [
    ("gpt-4.1", LiteLLMModelArgs, "openai/gpt-4.1-2025-04-14", chat_api_tools),
    # ("claude-3", LiteLLMModelArgs, "anthropic/claude-3-haiku-20240307", anthropic_tools),   # fails for parallel tool calls
    # ("claude-3.7", LiteLLMModelArgs, "anthropic/claude-3-7-sonnet-20250219", anthropic_tools), # fails for parallel tool calls
    ("claude-4-sonnet", LiteLLMModelArgs, "anthropic/claude-sonnet-4-20250514", chat_api_tools),
    # ("gpt-o3", LiteLLMModelArgs, "openai/o3-2025-04-16", chat_api_tools), # fails for parallel tool calls
    # add more models as needed
]

def add_user_messages(msg_builder):
    return [
        msg_builder.user().add_text("What is the weather in Paris and Delhi?"),
        msg_builder.user().add_text("You must call multiple tools to achieve the task."),
    ]

## Test multiaction
@pytest.mark.pricy
def test_multi_action_tool_calls():
    """
    Test that the model can produce multiple tool calls in parallel.
    Note: Remove assert and Uncomment commented lines to see the full behaviour of models and tool choices.
    """
    res_df = []
    for tool_choice in [
        # "none",
        "required",  # fails for Responses API
        "any",  # fails for Responses API
        "auto",
        # "get_weather",  # force a specific tool call
    ]:
        for name, llm_class, checkpoint_name, tools in tool_test_configs:
            model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
            llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
            messages = add_user_messages(msg_builder)
            if tool_choice == "get_weather":  # force a specific tool call
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, force_call_tool=tool_choice)
                )
            else:
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, tool_choice=tool_choice)
                )
            num_tool_calls = len(response.tool_calls) if response.tool_calls else 0
            row = {
                "model": name,
                "checkpoint": checkpoint_name,
                "tool_choice": tool_choice,
                "num_tool_calls": num_tool_calls,
                "action": response.action,
            }
            res_df.append(row)
            assert (
                num_tool_calls == 2
            ), f"Expected 2 tool calls, but got {num_tool_calls} for {name} with tool choice {tool_choice}"
    # import pandas as pd
    # print(pd.DataFrame(res_df))

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_single_tool_call():
    """
    Test that the LLMOutput contains only one tool call when use_only_first_toolcall is True.
    """
    for tool_choice in [
        # 'none',
        "required",
        "any",
        "auto",
    ]:
        for name, llm_class, checkpoint_name, tools in tool_test_configs:
            print(name, "tool choice:", tool_choice, "\n", "**" * 10)
            llm_class = partial(llm_class, use_only_first_toolcall=True)
            model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
            llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
            messages = add_user_messages(msg_builder)
            if tool_choice == "get_weather":  # force a specific tool call
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, force_call_tool=tool_choice)
                )
            else:
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, tool_choice=tool_choice)
                )
            num_tool_calls = len(response.tool_calls) if response.tool_calls else 0
            assert (
                num_tool_calls == 1
            ), f"Expected 1 tool calls, but got {num_tool_calls} for {name} with tool choice {tool_choice }"

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_force_tool_call():
    """
    Test that the model can produce a specific tool call when requested.
    The user message asks the 'weather' but we force call tool "get_time".
    We test if 'get_time' is present in the tool calls.
    Note: Model can have other tool calls as well.
    """
    force_call_tool = "get_time"
    for name, llm_class, checkpoint_name, tools in tool_test_configs:
        model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
        llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
        messages = add_user_messages(msg_builder)  # asks weather in Paris and Delhi
        response: LLMOutput = llm(
            APIPayload(messages=messages, tools=tools, force_call_tool=force_call_tool)
        )
        called_fn_names = [call.name for call in response.tool_calls] if response.tool_calls else []
        assert response.tool_calls is not None
        assert any(
            fn_name == "get_time" for fn_name in called_fn_names
        ), f"Model:{name},Expected all tool calls to be 'get_time', but got {called_fn_names} with force call {force_call_tool}"

if __name__ == "__main__":
    test_multi_action_tool_calls()
    test_force_tool_call()
    test_single_tool_call()


                    "description": "The location to get the time for.",
                }
            },
            "required": ["location"],
        },
    },
]

# test_config (setting name, BaseModelArgs, model_name, tools)
tool_test_configs = [
    ("gpt-4.1", LiteLLMModelArgs, "openai/gpt-4.1-2025-04-14", chat_api_tools),
    # ("claude-3", LiteLLMModelArgs, "anthropic/claude-3-haiku-20240307", anthropic_tools),   # fails for parallel tool calls
    # ("claude-3.7", LiteLLMModelArgs, "anthropic/claude-3-7-sonnet-20250219", anthropic_tools), # fails for parallel tool calls
    ("claude-4-sonnet", LiteLLMModelArgs, "anthropic/claude-sonnet-4-20250514", chat_api_tools),
    # ("gpt-o3", LiteLLMModelArgs, "openai/o3-2025-04-16", chat_api_tools), # fails for parallel tool calls
    # add more models as needed
]

def add_user_messages(msg_builder):
    return [
        msg_builder.user().add_text("What is the weather in Paris and Delhi?"),
        msg_builder.user().add_text("You must call multiple tools to achieve the task."),
    ]

## Test multiaction
@pytest.mark.pricy
def test_multi_action_tool_calls():
    """
    Test that the model can produce multiple tool calls in parallel.
    Note: Remove assert and Uncomment commented lines to see the full behaviour of models and tool choices.
    """
    res_df = []
    for tool_choice in [
        # "none",
        "required",  # fails for Responses API
        "any",  # fails for Responses API
        "auto",
        # "get_weather",  # force a specific tool call
    ]:
        for name, llm_class, checkpoint_name, tools in tool_test_configs:
            model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
            llm, msg_builder = model_args.make_model(), model_args.get_message_builder()


# test_config (setting name, BaseModelArgs, model_name, tools)
tool_test_configs = [
    ("gpt-4.1", LiteLLMModelArgs, "openai/gpt-4.1-2025-04-14", chat_api_tools),
    # ("claude-3", LiteLLMModelArgs, "anthropic/claude-3-haiku-20240307", anthropic_tools),   # fails for parallel tool calls
    # ("claude-3.7", LiteLLMModelArgs, "anthropic/claude-3-7-sonnet-20250219", anthropic_tools), # fails for parallel tool calls
    ("claude-4-sonnet", LiteLLMModelArgs, "anthropic/claude-sonnet-4-20250514", chat_api_tools),
    # ("gpt-o3", LiteLLMModelArgs, "openai/o3-2025-04-16", chat_api_tools), # fails for parallel tool calls
    # add more models as needed
]

def add_user_messages(msg_builder):
    return [
        msg_builder.user().add_text("What is the weather in Paris and Delhi?"),
        msg_builder.user().add_text("You must call multiple tools to achieve the task."),
    ]

## Test multiaction
@pytest.mark.pricy
def test_multi_action_tool_calls():
    """
    Test that the model can produce multiple tool calls in parallel.
    Note: Remove assert and Uncomment commented lines to see the full behaviour of models and tool choices.
    """
    res_df = []
    for tool_choice in [
        # "none",
        "required",  # fails for Responses API
        "any",  # fails for Responses API
        "auto",
        # "get_weather",  # force a specific tool call
    ]:
        for name, llm_class, checkpoint_name, tools in tool_test_configs:
            model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
            llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
            messages = add_user_messages(msg_builder)
            if tool_choice == "get_weather":  # force a specific tool call
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, force_call_tool=tool_choice)
                )
            else:
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, tool_choice=tool_choice)
                )
            num_tool_calls = len(response.tool_calls) if response.tool_calls else 0
            row = {
                "model": name,
                "checkpoint": checkpoint_name,
                "tool_choice": tool_choice,
                "num_tool_calls": num_tool_calls,
                "action": response.action,
            }
            res_df.append(row)
            assert (
                num_tool_calls == 2
            ), f"Expected 2 tool calls, but got {num_tool_calls} for {name} with tool choice {tool_choice}"
    # import pandas as pd
    # print(pd.DataFrame(res_df))

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_single_tool_call():
    """
    Test that the LLMOutput contains only one tool call when use_only_first_toolcall is True.
    """
    for tool_choice in [
        # 'none',
        "required",
        "any",
        "auto",
    ]:
        for name, llm_class, checkpoint_name, tools in tool_test_configs:
            print(name, "tool choice:", tool_choice, "\n", "**" * 10)
            llm_class = partial(llm_class, use_only_first_toolcall=True)
            model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)


                    APIPayload(messages=messages, tools=tools, tool_choice=tool_choice)
                )
            num_tool_calls = len(response.tool_calls) if response.tool_calls else 0
            row = {
                "model": name,
                "checkpoint": checkpoint_name,
                "tool_choice": tool_choice,
                "num_tool_calls": num_tool_calls,
                "action": response.action,
            }
            res_df.append(row)
            assert (
                num_tool_calls == 2
            ), f"Expected 2 tool calls, but got {num_tool_calls} for {name} with tool choice {tool_choice}"
    # import pandas as pd
    # print(pd.DataFrame(res_df))

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_single_tool_call():
    """
    Test that the LLMOutput contains only one tool call when use_only_first_toolcall is True.
    """
    for tool_choice in [
        # 'none',
        "required",
        "any",
        "auto",
    ]:
        for name, llm_class, checkpoint_name, tools in tool_test_configs:
            print(name, "tool choice:", tool_choice, "\n", "**" * 10)
            llm_class = partial(llm_class, use_only_first_toolcall=True)
            model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
            llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
            messages = add_user_messages(msg_builder)
            if tool_choice == "get_weather":  # force a specific tool call
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, force_call_tool=tool_choice)
                )
            else:
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, tool_choice=tool_choice)
                )
            num_tool_calls = len(response.tool_calls) if response.tool_calls else 0
            assert (
                num_tool_calls == 1
            ), f"Expected 1 tool calls, but got {num_tool_calls} for {name} with tool choice {tool_choice }"

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_force_tool_call():
    """
    Test that the model can produce a specific tool call when requested.
    The user message asks the 'weather' but we force call tool "get_time".
    We test if 'get_time' is present in the tool calls.
    Note: Model can have other tool calls as well.
    """
    force_call_tool = "get_time"
    for name, llm_class, checkpoint_name, tools in tool_test_configs:
        model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
        llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
        messages = add_user_messages(msg_builder)  # asks weather in Paris and Delhi
        response: LLMOutput = llm(
            APIPayload(messages=messages, tools=tools, force_call_tool=force_call_tool)
        )
        called_fn_names = [call.name for call in response.tool_calls] if response.tool_calls else []


            llm_class = partial(llm_class, use_only_first_toolcall=True)
            model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
            llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
            messages = add_user_messages(msg_builder)
            if tool_choice == "get_weather":  # force a specific tool call
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, force_call_tool=tool_choice)
                )
            else:
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, tool_choice=tool_choice)
                )
            num_tool_calls = len(response.tool_calls) if response.tool_calls else 0
            assert (
                num_tool_calls == 1
            ), f"Expected 1 tool calls, but got {num_tool_calls} for {name} with tool choice {tool_choice }"

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_force_tool_call():
    """
    Test that the model can produce a specific tool call when requested.
    The user message asks the 'weather' but we force call tool "get_time".
    We test if 'get_time' is present in the tool calls.
    Note: Model can have other tool calls as well.
    """
    force_call_tool = "get_time"
    for name, llm_class, checkpoint_name, tools in tool_test_configs:
        model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
        llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
        messages = add_user_messages(msg_builder)  # asks weather in Paris and Delhi
        response: LLMOutput = llm(
            APIPayload(messages=messages, tools=tools, force_call_tool=force_call_tool)
        )
        called_fn_names = [call.name for call in response.tool_calls] if response.tool_calls else []
        assert response.tool_calls is not None
        assert any(
            fn_name == "get_time" for fn_name in called_fn_names
        ), f"Model:{name},Expected all tool calls to be 'get_time', but got {called_fn_names} with force call {force_call_tool}"

if __name__ == "__main__":
    test_multi_action_tool_calls()
    test_force_tool_call()
    test_single_tool_call()


import os
from typing import Any, Dict, List, Optional
from unittest.mock import MagicMock, patch

import anthropic
import openai
import pytest

from agentlab.llm import tracking
from agentlab.llm.response_api import (
    AnthropicAPIMessageBuilder,
    APIPayload,
    ClaudeResponseModelArgs,
    LLMOutput,
    OpenAIChatCompletionAPIMessageBuilder,
    OpenAIChatModelArgs,
    OpenAIResponseAPIMessageBuilder,
    OpenAIResponseModelArgs,
)

# Helper to create a mock OpenAI ChatCompletion response
def create_mock_openai_chat_completion(
    content=None, tool_calls=None, prompt_tokens=10, completion_tokens=20
):
    completion = MagicMock(spec=openai.types.chat.ChatCompletion)
    choice = MagicMock()
    message = MagicMock(spec=openai.types.chat.ChatCompletionMessage)
    message.content = content
    message.tool_calls = None
    if tool_calls:
        message.tool_calls = []
        for tc in tool_calls:
            tool_call_mock = MagicMock(
                spec=openai.types.chat.chat_completion_message_tool_call.ChatCompletionMessageToolCall
            )
            tool_call_mock.id = tc["id"]
            tool_call_mock.type = tc["type"]
            tool_call_mock.function = MagicMock()
            tool_call_mock.function.name = tc["function"]["name"]
            tool_call_mock.function.arguments = tc["function"]["arguments"]
            message.tool_calls.append(tool_call_mock)

    choice.message = message
    completion.choices = [choice]

    completion.usage = MagicMock()
    # Explicitly set the attributes that get_tokens_counts_from_response will try first.
    # These are the generic names.
    completion.usage.input_tokens = prompt_tokens
    completion.usage.output_tokens = completion_tokens

    # Also set the OpenAI-specific names if any other part of the code might look for them directly,
    # or if get_tokens_counts_from_response had different fallback logic.
    completion.usage.prompt_tokens = prompt_tokens
    completion.usage.completion_tokens = completion_tokens
    prompt_tokens_details_mock = MagicMock()
    prompt_tokens_details_mock.cached_tokens = 0
    completion.usage.prompt_tokens_details = prompt_tokens_details_mock

    completion.model_dump.return_value = {
        "id": "chatcmpl-xxxx",
        "choices": [
            {"message": {"role": "assistant", "content": content, "tool_calls": tool_calls}}
        ],
        # Ensure the usage dict in model_dump also reflects the token counts accurately.
        # The get_tokens_counts_from_response also has a path for dict style.
        "usage": {
            "input_tokens": prompt_tokens,  # Generic name
            "output_tokens": completion_tokens,  # Generic name
            "prompt_tokens": prompt_tokens,  # OpenAI specific
            "completion_tokens": completion_tokens,  # OpenAI specific
            "prompt_tokens_details": {"cached_tokens": 0},
        },
    }
    message.to_dict.return_value = {
        "role": "assistant",
        "content": content,
        "tool_calls": tool_calls,
    }
    return completion

responses_api_tools = [
    {
        "type": "function",
        "name": "get_weather",
        "description": "Get the current weather in a given location.",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The location to get the weather for.",
                },
                "unit": {
                    "type": "string",
                    "enum": ["celsius", "fahrenheit"],
                    "description": "The unit of temperature.",
                },
            },
            "required": ["location"],
        },
    }
]

chat_api_tools = [
    {
        "type": "function",
        "name": "get_weather",
        "description": "Get the current weather in a given location.",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The location to get the weather for.",
                },
                "unit": {
                    "type": "string",
                    "enum": ["celsius", "fahrenheit"],
                    "description": "The unit of temperature.",
                },
            },
            "required": ["location"],
        },
    }
]
anthropic_tools = [
    {
        "name": "get_weather",
        "description": "Get the current weather in a given location.",
        "input_schema": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The location to get the weather for.",
                },
            },
            "required": ["location"],
        },
    }
]

# Helper to create a mock Anthropic response
def create_mock_anthropic_response(
    text_content=None, tool_use=None, input_tokens=15, output_tokens=25
):

    response = MagicMock(spec=anthropic.types.Message)
    response.type = "message"  # Explicitly set the type attribute
    response.content = []
    response.content = []
    if text_content:
        text_block = MagicMock(spec=anthropic.types.TextBlock)
        text_block.type = "text"
        text_block.text = text_content
        response.content.append(text_block)
    if tool_use:
        tool_use_block = MagicMock(spec=anthropic.types.ToolUseBlock)
        tool_use_block.type = "tool_use"
        tool_use_block.id = tool_use["id"]
        tool_use_block.name = tool_use["name"]
        tool_use_block.input = tool_use["input"]
        response.content.append(tool_use_block)
    response.usage = MagicMock()
    response.usage.input_tokens = input_tokens
    response.usage.output_tokens = output_tokens
    response.usage.cache_input_tokens = 0
    response.usage.cache_creation_input_tokens = 0
    return response

def create_mock_openai_responses_api_response(
    outputs: Optional[List[Dict[str, Any]]] = None, input_tokens: int = 10, output_tokens: int = 20
) -> MagicMock:
    """
    Helper to create a mock response object similar to what
    openai.resources.Responses.create() would return.
    Compatible with OpenAIResponseModel and TrackAPIPricingMixin.
    """

    response_mock = MagicMock(spec=openai.types.responses.response.Response)
    response_mock.type = "response"
    response_mock.output = []

    if outputs:
        for out_data in outputs:
            output_item_mock = MagicMock()
            output_item_mock.type = out_data.get("type")

            if output_item_mock.type == "function_call":
                # You can adapt this depending on your expected object structure
                output_item_mock.name = out_data.get("name")
                output_item_mock.arguments = out_data.get("arguments")
                output_item_mock.call_id = out_data.get("call_id")
            elif output_item_mock.type == "reasoning":
                output_item_mock.summary = []
                for text_content in out_data.get("summary", []):
                    summary_text_mock = MagicMock()
                    summary_text_mock.text = text_content
                    output_item_mock.summary.append(summary_text_mock)

            response_mock.output.append(output_item_mock)

    # Token usage for pricing tracking
    response_mock.usage = MagicMock(spec=openai.types.responses.response.ResponseUsage)
    response_mock.usage.input_tokens = input_tokens
    response_mock.usage.output_tokens = output_tokens
    response_mock.usage.prompt_tokens = input_tokens
    response_mock.usage.completion_tokens = output_tokens
    input_tokens_details_mock = MagicMock()
    input_tokens_details_mock.cached_tokens = 0
    response_mock.usage.input_tokens_details = input_tokens_details_mock

    return response_mock

# --- Test MessageBuilders ---

def test_openai_response_api_message_builder_text():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_text("Hello, world!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "input_text", "text": "Hello, world!"}]

def test_openai_response_api_message_builder_image():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,SIMPLEBASE64STRING")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [
        {"type": "input_image", "image_url": "data:image/png;base64,SIMPLEBASE64STRING"}
    ]

def test_anthropic_api_message_builder_text():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_text("Hello, Anthropic!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, Anthropic!"}]

def test_anthropic_api_message_builder_image():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,ANTHROPICBASE64")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert len(messages[0]["content"]) == 1
    image_content = messages[0]["content"][0]
    assert image_content["type"] == "image"
    assert image_content["source"]["type"] == "base64"
    assert image_content["source"]["media_type"] == "image/png"
    assert image_content["source"]["data"] == "ANTHROPICBASE64"  # Base64 prefix should be stripped

def test_openai_chat_completion_api_message_builder_text():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_text("Hello, ChatCompletion!")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, ChatCompletion!"}]

def test_openai_chat_completion_api_message_builder_image():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_image("data:image/jpeg;base64,CHATCOMPLETIONBASE64")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [
        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,CHATCOMPLETIONBASE64"}}
    ]

def test_openai_chat_completion_model_parse_and_cost():
    args = OpenAIChatModelArgs(model_name="gpt-3.5-turbo")
    with patch("agentlab.llm.response_api.OpenAI") as mock_openai_class:
        mock_client = MagicMock()
        mock_openai_class.return_value = mock_client
        model = args.make_model()

    mock_response = create_mock_openai_chat_completion(
        content="This is a test thought.",
        tool_calls=[
            {
                "id": "call_123",
                "type": "function",
                "function": {"name": "get_weather", "arguments": '{"location": "Paris"}'},
            }
        ],
        prompt_tokens=50,
        completion_tokens=30,
    )

    with patch.object(
        model.client.chat.completions, "create", return_value=mock_response
    ) as mock_create:
        with tracking.set_tracker() as global_tracker:
            messages = [
                OpenAIChatCompletionAPIMessageBuilder.user().add_text(
                    "What's the weather in Paris?"
                )
            ]
            payload = APIPayload(messages=messages)
            parsed_output = model(payload)

    mock_create.assert_called_once()
    assert parsed_output.raw_response.choices[0].message.content == "This is a test thought."
    assert parsed_output.action == """get_weather(location='Paris')"""
    assert parsed_output.raw_response.choices[0].message.tool_calls[0].id == "call_123"
    # Check cost tracking (token counts)
    assert global_tracker.stats["input_tokens"] == 50
    assert global_tracker.stats["output_tokens"] == 30
    assert global_tracker.stats["cost"] > 0

def test_claude_response_model_parse_and_cost():
    args = ClaudeResponseModelArgs(model_name="claude-3-haiku-20240307")
    model = args.make_model()

    mock_anthropic_api_response = create_mock_anthropic_response(
        text_content="Thinking about the request.",
        tool_use={"id": "tool_abc", "name": "search_web", "input": {"query": "latest news"}},
        input_tokens=40,
        output_tokens=20,
    )

    with patch.object(
        model.client.messages, "create", return_value=mock_anthropic_api_response
    ) as mock_create:
        with tracking.set_tracker() as global_tracker:
            messages = [AnthropicAPIMessageBuilder.user().add_text("Search for latest news")]
            payload = APIPayload(messages=messages)
            parsed_output = model(payload)

    mock_create.assert_called_once()
    fn_call = next(iter(parsed_output.tool_calls))

    assert "Thinking about the request." in parsed_output.think
    assert parsed_output.action == """search_web(query='latest news')"""
    assert fn_call.name == "search_web"
    assert global_tracker.stats["input_tokens"] == 40
    assert global_tracker.stats["output_tokens"] == 20

def test_openai_response_model_parse_and_cost():
    args = OpenAIResponseModelArgs(model_name="gpt-4.1")

    mock_function_call_output = {
        "type": "function_call",
        "name": "get_current_weather",
        "arguments": '{"location": "Boston, MA", "unit": "celsius"}',
        "call_id": "call_abc123",
    }

    mock_api_resp = create_mock_openai_responses_api_response(
        outputs=[mock_function_call_output],
        input_tokens=70,
        output_tokens=40,
    )

    with patch("agentlab.llm.response_api.OpenAI") as mock_openai_class:
        mock_client = MagicMock()
        mock_openai_class.return_value = mock_client
        model = args.make_model()

    with patch.object(
        model.client.responses, "create", return_value=mock_api_resp
    ) as mock_create_method:
        with tracking.set_tracker() as global_tracker:
            messages = [
                OpenAIResponseAPIMessageBuilder.user().add_text("What's the weather in Boston?")
            ]
            payload = APIPayload(messages=messages)
            parsed_output = model(payload)

    mock_create_method.assert_called_once()
    fn_calls = [
        content
        for content in parsed_output.tool_calls.raw_calls.output
        if content.type == "function_call"
    ]
    assert parsed_output.action == "get_current_weather(location='Boston, MA', unit='celsius')"
    assert fn_calls[0].call_id == "call_abc123"
    assert parsed_output.raw_response == mock_api_resp
    assert global_tracker.stats["input_tokens"] == 70
    assert global_tracker.stats["output_tokens"] == 40

# --- Test Response Models (Pricy - require API keys and actual calls) ---

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_chat_completion_model_pricy_call():
    """Tests OpenAIChatCompletionModel with a real API call."""
    args = OpenAIChatModelArgs(
        model_name="gpt-4.1",
        temperature=1e-5,
        max_new_tokens=100,
    )

    tools = chat_api_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [
            OpenAIChatCompletionAPIMessageBuilder.user().add_text("What is the weather in Paris?")
        ]
        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed_output = model(payload)

    assert parsed_output.raw_response is not None
    assert (
        parsed_output.action == "get_weather(location='Paris')"
    ), f""" Expected get_weather(location='Paris') but got {parsed_output.action}"""
    assert global_tracker.stats["input_tokens"] > 0
    assert global_tracker.stats["output_tokens"] > 0
    assert global_tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
def test_claude_response_model_pricy_call():
    """Tests ClaudeResponseModel with a real API call."""

    args = ClaudeResponseModelArgs(
        model_name="claude-3-haiku-20240307",
        temperature=1e-5,
        max_new_tokens=100,
    )
    tools = anthropic_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [AnthropicAPIMessageBuilder.user().add_text("What is the weather in Paris?")]
        payload = APIPayload(messages=messages, tools=tools)
        parsed_output = model(payload)

    assert parsed_output.raw_response is not None
    assert (
        parsed_output.action == "get_weather(location='Paris')"
    ), f"""Expected get_weather('Paris') but got {parsed_output.action}"""
    assert global_tracker.stats["input_tokens"] > 0
    assert global_tracker.stats["output_tokens"] > 0
    assert global_tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_response_model_pricy_call():
    """
    Tests OpenAIResponseModel output parsing and cost tracking with both
    function_call and reasoning outputs.
    """
    args = OpenAIResponseModelArgs(model_name="gpt-4.1", temperature=1e-5, max_new_tokens=100)

    tools = responses_api_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [
            OpenAIResponseAPIMessageBuilder.user().add_text("What is the weather in Paris?")
        ]
        payload = APIPayload(messages=messages, tools=tools)
        parsed_output = model(payload)

    assert parsed_output.raw_response is not None
    assert (
        parsed_output.action == """get_weather(location='Paris', unit='celsius')"""
    ), f""" Expected get_weather(location='Paris', unit='celsius') but got {parsed_output.action}"""
    assert global_tracker.stats["input_tokens"] > 0
    assert global_tracker.stats["output_tokens"] > 0
    assert global_tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_response_model_with_multiple_messages_and_cost_tracking():
    """
    Test OpenAIResponseModel's output parsing and cost tracking
    with a tool-using assistant and follow-up interaction.
    """
    args = OpenAIResponseModelArgs(model_name="gpt-4.1", temperature=1e-5, max_new_tokens=100)

    tools = responses_api_tools
    model = args.make_model()
    builder = args.get_message_builder()

    messages = [builder.user().add_text("What is the weather in Paris?")]

    with tracking.set_tracker() as tracker:
        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed = model(payload)
        prev_input = tracker.stats["input_tokens"]
        prev_output = tracker.stats["output_tokens"]
        prev_cost = tracker.stats["cost"]

        assert parsed.tool_calls, "Expected tool calls in the response"
        # Set tool responses
        for tool_call in parsed.tool_calls:
            tool_call.response_text("Its sunny! 25°C")
        # Simulate tool execution and user follow-up
        messages += [
            builder.add_responded_tool_calls(parsed.tool_calls),
            builder.user().add_text("What is the weather in Delhi?"),
        ]

        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed = model(payload)

        delta_input = tracker.stats["input_tokens"] - prev_input
        delta_output = tracker.stats["output_tokens"] - prev_output
        delta_cost = tracker.stats["cost"] - prev_cost

    assert prev_input > 0
    assert prev_output > 0
    assert prev_cost > 0
    assert parsed.raw_response is not None
    assert (
        parsed.action == """get_weather(location='Delhi', unit='celsius')"""
    ), f
# ... truncated ...


from unittest.mock import MagicMock, patch

import anthropic
import openai
import pytest

from agentlab.llm import tracking
from agentlab.llm.response_api import (
    AnthropicAPIMessageBuilder,
    APIPayload,
    ClaudeResponseModelArgs,
    LLMOutput,
    OpenAIChatCompletionAPIMessageBuilder,
    OpenAIChatModelArgs,
    OpenAIResponseAPIMessageBuilder,
    OpenAIResponseModelArgs,
)

# Helper to create a mock OpenAI ChatCompletion response
def create_mock_openai_chat_completion(
    content=None, tool_calls=None, prompt_tokens=10, completion_tokens=20
):
    completion = MagicMock(spec=openai.types.chat.ChatCompletion)
    choice = MagicMock()
    message = MagicMock(spec=openai.types.chat.ChatCompletionMessage)
    message.content = content
    message.tool_calls = None
    if tool_calls:
        message.tool_calls = []
        for tc in tool_calls:
            tool_call_mock = MagicMock(
                spec=openai.types.chat.chat_completion_message_tool_call.ChatCompletionMessageToolCall
            )
            tool_call_mock.id = tc["id"]
            tool_call_mock.type = tc["type"]
            tool_call_mock.function = MagicMock()
            tool_call_mock.function.name = tc["function"]["name"]
            tool_call_mock.function.arguments = tc["function"]["arguments"]
            message.tool_calls.append(tool_call_mock)

    choice.message = message
    completion.choices = [choice]

    completion.usage = MagicMock()
    # Explicitly set the attributes that get_tokens_counts_from_response will try first.
    # These are the generic names.
    completion.usage.input_tokens = prompt_tokens
    completion.usage.output_tokens = completion_tokens

    # Also set the OpenAI-specific names if any other part of the code might look for them directly,
    # or if get_tokens_counts_from_response had different fallback logic.
    completion.usage.prompt_tokens = prompt_tokens
    completion.usage.completion_tokens = completion_tokens
    prompt_tokens_details_mock = MagicMock()
    prompt_tokens_details_mock.cached_tokens = 0
    completion.usage.prompt_tokens_details = prompt_tokens_details_mock

    completion.model_dump.return_value = {
        "id": "chatcmpl-xxxx",
        "choices": [
            {"message": {"role": "assistant", "content": content, "tool_calls": tool_calls}}
        ],
        # Ensure the usage dict in model_dump also reflects the token counts accurately.
        # The get_tokens_counts_from_response also has a path for dict style.
        "usage": {
            "input_tokens": prompt_tokens,  # Generic name
            "output_tokens": completion_tokens,  # Generic name
            "prompt_tokens": prompt_tokens,  # OpenAI specific
            "completion_tokens": completion_tokens,  # OpenAI specific
            "prompt_tokens_details": {"cached_tokens": 0},
        },
    }
    message.to_dict.return_value = {
        "role": "assistant",
        "content": content,
        "tool_calls": tool_calls,
    }
    return completion

responses_api_tools = [
    {
        "type": "function",
        "name": "get_weather",
        "description": "Get the current weather in a given location.",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The location to get the weather for.",
                },
                "unit": {
                    "type": "string",
                    "enum": ["celsius", "fahrenheit"],
                    "description": "The unit of temperature.",
                },
            },


]
anthropic_tools = [
    {
        "name": "get_weather",
        "description": "Get the current weather in a given location.",
        "input_schema": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The location to get the weather for.",
                },
            },
            "required": ["location"],
        },
    }
]

# Helper to create a mock Anthropic response
def create_mock_anthropic_response(
    text_content=None, tool_use=None, input_tokens=15, output_tokens=25
):

    response = MagicMock(spec=anthropic.types.Message)
    response.type = "message"  # Explicitly set the type attribute
    response.content = []
    response.content = []
    if text_content:
        text_block = MagicMock(spec=anthropic.types.TextBlock)
        text_block.type = "text"
        text_block.text = text_content
        response.content.append(text_block)
    if tool_use:
        tool_use_block = MagicMock(spec=anthropic.types.ToolUseBlock)
        tool_use_block.type = "tool_use"
        tool_use_block.id = tool_use["id"]
        tool_use_block.name = tool_use["name"]
        tool_use_block.input = tool_use["input"]
        response.content.append(tool_use_block)
    response.usage = MagicMock()
    response.usage.input_tokens = input_tokens
    response.usage.output_tokens = output_tokens
    response.usage.cache_input_tokens = 0
    response.usage.cache_creation_input_tokens = 0
    return response

def create_mock_openai_responses_api_response(
    outputs: Optional[List[Dict[str, Any]]] = None, input_tokens: int = 10, output_tokens: int = 20
) -> MagicMock:
    """
    Helper to create a mock response object similar to what
    openai.resources.Responses.create() would return.
    Compatible with OpenAIResponseModel and TrackAPIPricingMixin.
    """

    response_mock = MagicMock(spec=openai.types.responses.response.Response)
    response_mock.type = "response"
    response_mock.output = []

    if outputs:
        for out_data in outputs:
            output_item_mock = MagicMock()
            output_item_mock.type = out_data.get("type")


    if text_content:
        text_block = MagicMock(spec=anthropic.types.TextBlock)
        text_block.type = "text"
        text_block.text = text_content
        response.content.append(text_block)
    if tool_use:
        tool_use_block = MagicMock(spec=anthropic.types.ToolUseBlock)
        tool_use_block.type = "tool_use"
        tool_use_block.id = tool_use["id"]
        tool_use_block.name = tool_use["name"]
        tool_use_block.input = tool_use["input"]
        response.content.append(tool_use_block)
    response.usage = MagicMock()
    response.usage.input_tokens = input_tokens
    response.usage.output_tokens = output_tokens
    response.usage.cache_input_tokens = 0
    response.usage.cache_creation_input_tokens = 0
    return response

def create_mock_openai_responses_api_response(
    outputs: Optional[List[Dict[str, Any]]] = None, input_tokens: int = 10, output_tokens: int = 20
) -> MagicMock:
    """
    Helper to create a mock response object similar to what
    openai.resources.Responses.create() would return.
    Compatible with OpenAIResponseModel and TrackAPIPricingMixin.
    """

    response_mock = MagicMock(spec=openai.types.responses.response.Response)
    response_mock.type = "response"
    response_mock.output = []

    if outputs:
        for out_data in outputs:
            output_item_mock = MagicMock()
            output_item_mock.type = out_data.get("type")

            if output_item_mock.type == "function_call":
                # You can adapt this depending on your expected object structure
                output_item_mock.name = out_data.get("name")
                output_item_mock.arguments = out_data.get("arguments")
                output_item_mock.call_id = out_data.get("call_id")
            elif output_item_mock.type == "reasoning":
                output_item_mock.summary = []
                for text_content in out_data.get("summary", []):
                    summary_text_mock = MagicMock()
                    summary_text_mock.text = text_content
                    output_item_mock.summary.append(summary_text_mock)

            response_mock.output.append(output_item_mock)

    # Token usage for pricing tracking
    response_mock.usage = MagicMock(spec=openai.types.responses.response.ResponseUsage)
    response_mock.usage.input_tokens = input_tokens
    response_mock.usage.output_tokens = output_tokens
    response_mock.usage.prompt_tokens = input_tokens
    response_mock.usage.completion_tokens = output_tokens
    input_tokens_details_mock = MagicMock()
    input_tokens_details_mock.cached_tokens = 0
    response_mock.usage.input_tokens_details = input_tokens_details_mock

    return response_mock

# --- Test MessageBuilders ---

def test_openai_response_api_message_builder_text():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_text("Hello, world!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "input_text", "text": "Hello, world!"}]

def test_openai_response_api_message_builder_image():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,SIMPLEBASE64STRING")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"


                    output_item_mock.summary.append(summary_text_mock)

            response_mock.output.append(output_item_mock)

    # Token usage for pricing tracking
    response_mock.usage = MagicMock(spec=openai.types.responses.response.ResponseUsage)
    response_mock.usage.input_tokens = input_tokens
    response_mock.usage.output_tokens = output_tokens
    response_mock.usage.prompt_tokens = input_tokens
    response_mock.usage.completion_tokens = output_tokens
    input_tokens_details_mock = MagicMock()
    input_tokens_details_mock.cached_tokens = 0
    response_mock.usage.input_tokens_details = input_tokens_details_mock

    return response_mock

# --- Test MessageBuilders ---

def test_openai_response_api_message_builder_text():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_text("Hello, world!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "input_text", "text": "Hello, world!"}]

def test_openai_response_api_message_builder_image():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,SIMPLEBASE64STRING")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [
        {"type": "input_image", "image_url": "data:image/png;base64,SIMPLEBASE64STRING"}
    ]

def test_anthropic_api_message_builder_text():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_text("Hello, Anthropic!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, Anthropic!"}]