import gradio as gr
from autogen import UserProxyAgent, ConversableAgent, register_function

def run_flow(prompt: str, history) -> None:
    flow = "default"

    llm_config_gpt4 = {
        "model": "gpt-4",
        "temperature": 0.7,
        # "max_tokens": 2048,
        "api_key" : "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",
    }

    user_proxy = ConversableAgent(
        name="User",
        llm_config=False,
        is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
        human_input_mode="NEVER",
    )

    assistant = ConversableAgent(
        name="Assistant",
        system_message="You are a helpful AI assistant. "
        "Return 'TERMINATE' when the task is done.",
        llm_config=llm_config_gpt4,
    )
    
    user_proxy.initiate_chat(
        assistant,
        message=prompt,
    )

    messages = user_proxy.chat_messages[assistant]
    print(messages)
    return messages
    

app = gr.ChatInterface(run_flow)
app.launch()