import autogen
import tempfile

# direct access to Ollama since 0.1.24, compatible with OpenAI /chat/completions
BASE_URL="http://localhost:11434/v1"

config_list_core = [
    {
        'base_url': BASE_URL,
        'api_key': "fakekey",
        'model': "llama3:latest",
    }
]

config_list_coder = [
    {
        'base_url': BASE_URL,
        'api_key': "fakekey",
        'model': "dolphin-llama3:latest",
    }
]

llm_config_core={
    "config_list": config_list_core,
}

llm_config_code={
    "config_list": config_list_coder,
}

use_groupchat = False

user_proxy = autogen.UserProxyAgent(
    name="user_proxy",
    human_input_mode="NEVER",
    #human_input_mode="TERMINATE",
    max_consecutive_auto_reply=10,
    is_termination_msg=lambda x: x.get("content", "").rstrip().endswith("TERMINATE"),
    code_execution_config={"work_dir": "coding", "use_docker":False},
    llm_config=llm_config_core,
    system_message="""Reply TERMINATE if the task has been solved at full satisfaction.
Otherwise, reply CONTINUE, or the reason why the task is not solved yet."""
)

task="""
Write a python script to output numbers 1 to 100 and then the user_proxy agent should run the script
"""


# Create a temporary directory
with tempfile.TemporaryDirectory() as temp_dir:
    print(f"Created temporary directory: {temp_dir}")
    

    # The temporary directory and its contents are automatically cleaned up
    # when the 'with' block is exited


    assistant = autogen.AssistantAgent(
        name="Assistant",
        llm_config=llm_config_core,
        # code_execution=False # Disable code execution entirely
        code_execution_config={"work_dir":temp_dir, "use_docker":False}
    )

    coder = autogen.AssistantAgent(
        name="Coder",
        llm_config=llm_config_code,
        # code_execution=False # Disable code execution entirely
        code_execution_config={"work_dir":temp_dir, "use_docker":False}
    )

    use_groupchat = True
    if use_groupchat:
        groupchat = autogen.GroupChat(agents=[user_proxy, coder, assistant], messages=[], max_round=12)
        manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config_core)
        user_proxy.initiate_chat(manager, message=task)
    else:
        user_proxy.initiate_chat(coder, message=task)