# Chat between two comedian agents

# 1. Import our agent class
from autogen import ConversableAgent, LLMConfig
from autogen.io.run_response import Cost, RunResponseProtocol

# 2. Define our LLM configuration for OpenAI's GPT-4o mini,
#    uses the OPENAI_API_KEY environment variable
# llm_config = LLMConfig(api_type="openai", model="gpt-4o-mini")
llm_config = LLMConfig.from_json(path="OAI_CONFIG_LIST").where(model="gpt-4o-mini")
print(f"Using LLM: {llm_config}")

# 3. Create our agents who will tell each other jokes,
#    with Jack ending the chat when Emma says FINISH
with llm_config:
    jack = ConversableAgent(
        "Jack",
        system_message=("Your name is Jack and you are a comedian in a two-person comedy show."),
        is_termination_msg=lambda x: "FINISH" in x["content"],
    )
    emma = ConversableAgent(
        "Emma",
        system_message=(
            "Your name is Emma and you are a comedian "
            "in a two-person comedy show. Say the word FINISH "
            "ONLY AFTER you've heard 2 of Jack's jokes."
        ),
    )

# 4. Run the chat
response: RunResponseProtocol = jack.run(
    emma, message="Emma, tell me a joke about goldfish and peanut butter.", summary_method="reflection_with_llm"
)

for event in response.events:
    print(event)

    if event.type == "input_request":
        event.content.respond("exit")

print(f"{response.summary=}")
print(f"{response.messages=}")
print(f"{response.events=}")
print(f"{response.context_variables=}")
print(f"{response.last_speaker=}")
print(f"{response.cost=}")
assert response.last_speaker in ["Jack", "Emma"], "Last speaker should be one of the agents"
assert isinstance(response.cost, Cost)

# Chat between two comedian agents

# 1. Import console event processor

# 2. Create our agents who will tell each other jokes,
#    with Jack ending the chat when Emma says FINISH
with llm_config:
    jack = ConversableAgent(
        "Jack",
        system_message=("Your name is Jack and you are a comedian in a two-person comedy show."),
        is_termination_msg=lambda x: "FINISH" in x["content"],
        human_input_mode="NEVER",
    )
    emma = ConversableAgent(
        "Emma",
        system_message=(
            "Your name is Emma and you are a comedian "
            "in a two-person comedy show. Say the word FINISH "
            "ONLY AFTER you've heard 2 of Jack's jokes."
        ),
        human_input_mode="NEVER",
    )

# 3. Run the chat
response = jack.run(
    emma, message="Emma, tell me a joke about goldfish and peanut butter.", summary_method="reflection_with_llm"
)

response.process()

assert response.last_speaker in ["Jack", "Emma"], "Last speaker should be one of the agents"
assert response.summary is not None, "Summary should not be None"
assert len(response.messages) > 0, "Messages should not be empty"

# 1. Create our LLM agent
with llm_config:
    my_agent = ConversableAgent(
        name="helpful_agent",
        system_message="You are a poetic AI assistant, respond in rhyme.",
    )

# 2. Run the agent with a prompt
response = my_agent.run(
    message="In one sentence, what's the big deal about AI?", max_turns=1, summary_method="reflection_with_llm"
)

response.process()
print(f"{response.summary=}")
print(f"{response.messages=}")
print(f"{response.last_speaker=}")

assert response.summary is not None, "Summary should not be None"
assert len(response.messages) == 2, "Messages should not be empty"
assert response.last_speaker == "helpful_agent", "Last speaker should be an agent"

# Group chat amongst agents to create a 4th grade lesson plan
# Flow determined by Group Chat Manager automatically, and
# should be Teacher > Planner > Reviewer > Teacher (repeats if necessary)

# 1. Import our agent and group chat classes
from autogen import GroupChat, GroupChatManager

with llm_config:
    # Planner agent setup
    planner_message = "Create lesson plans for 4th grade. Use format: <title>, <learning_objectives>, <script>"
    planner = ConversableAgent(name="planner_agent", system_message=planner_message, description="Creates lesson plans")

    # Reviewer agent setup
    reviewer_message = "Review lesson plans against 4th grade curriculum. Provide max 3 changes."
    reviewer = ConversableAgent(
        name="reviewer_agent", system_message=reviewer_message, description="Reviews lesson plans"
    )

    # Teacher agent setup
    teacher_message = "Choose topics and work with planner and reviewer. Say DONE! when finished."
    teacher = ConversableAgent(
        name="teacher_agent",
        system_message=teacher_message,
    )

# Setup group chat
groupchat = GroupChat(agents=[teacher, planner, reviewer], speaker_selection_method="auto", messages=[])

# Create manager
# At each turn, the manager will check if the message contains DONE! and end the chat if so
# Otherwise, it will select the next appropriate agent using its LLM
manager = GroupChatManager(
    name="group_manager",
    groupchat=groupchat,
    llm_config=llm_config,
    is_termination_msg=lambda x: "DONE!" in (x.get("content", "") or "").upper(),
)

# Start the conversation
response = teacher.run(
    recipient=manager, message="Let's teach the kids about the solar system.", summary_method="reflection_with_llm"
)

response.process()

print(f"{response.summary=}")
print(f"{response.messages=}")
print(f"{response.last_speaker=}")

assert response.summary is not None, "Summary should not be None"
assert len(response.messages) > 0, "Messages should not be empty"
assert response.last_speaker in ["teacher_agent", "planner_agent", "reviewer_agent"], (
    "Last speaker should be one of the agents"
)

from autogen import run_swarm
from autogen.agentchat.contrib.swarm_agent import AfterWorkOption

# 1. Create our agents
planner_message = """You are a classroom lesson planner.
Given a topic, write a lesson plan for a fourth grade class.
If you are given revision feedback, update your lesson plan and record it.
Use the following format:
<title>Lesson plan title</title>
<learning_objectives>Key learning objectives</learning_objectives>
<script>How to introduce the topic to the kids</script>
"""

reviewer_message = """You are a classroom lesson reviewer.
You compare the lesson plan to the fourth grade curriculum
and provide a maximum of 3 recommended changes for each review.
Make sure you provide recommendations each time the plan is updated.
"""

teacher_message = """You are a classroom teacher.
You decide topics for lessons and work with a lesson planner.
and reviewer to create and finalise lesson plans.
"""

with llm_config:
    lesson_planner = ConversableAgent(name="planner_agent", system_message=planner_message)

    lesson_reviewer = ConversableAgent(name="reviewer_agent", system_message=reviewer_message)

    teacher = ConversableAgent(
        name="teacher_agent",
        system_message=teacher_message,
    )

# 2. Initiate the swarm chat using a swarm manager who will
# select agents automatically
response = run_swarm(
    initial_agent=teacher,
    agents=[lesson_planner, lesson_reviewer, teacher],
    messages="Today, let's introduce our kids to the solar system.",
    max_rounds=10,
    swarm_manager_args={"llm_config": llm_config},
    after_work=AfterWorkOption.SWARM_MANAGER,
)

response.process()

# for events in response.events:
#     if events.type == "input_request":
#         events.content.respond("exit")

print(f"{response.summary=}")
print(f"{response.messages=}")
print(f"{response.last_speaker=}")

assert response.summary is not None, "Summary should not be None"
assert len(response.messages) > 0, "Messages should not be empty"
assert response.last_speaker in ["teacher_agent", "planner_agent", "reviewer_agent"], (
    "Last speaker should be one of the agents"
)

from autogen.agentchat.user_proxy_agent import UserProxyAgent

financial_tasks = [
    """What are the current stock prices of NVDA and TESLA, and how is the performance over the past month in terms of percentage change?""",
    """Investigate possible reasons of the stock performance.""",
]

writing_tasks = ["""Develop an engaging blog post using any information provided."""]

with llm_config:
    financial_assistant = ConversableAgent(
        name="Financial_assistant",
        system_message="You are a financial assistant, helping with stock market analysis. Reply 'TERMINATE' when financial tasks are done.",
        llm_config=llm_config,
    )
    research_assistant = ConversableAgent(
        name="Researcher",
        system_message="You are a research assistant, helping with stock market analysis. Reply 'TERMINATE' when research tasks are done.",
        llm_config=llm_config,
    )
    writer = ConversableAgent(
        name="writer",
        llm_config=llm_config,
        system_message="""
            You are a professional writer, known for
            your insightful and engaging articles.
            You transform complex concepts into compelling narratives.
            Reply "TERMINATE" in the end when everything is done.
            """,
    )

    user = UserProxyAgent(
        name="User",
        human_input_mode="NEVER",
        is_termination_msg=lambda x: x.get("content", "") and x.get("content", "").rstrip().endswith("TERMINATE"),
        code_execution_config={
            "last_n_messages": 1,
            "work_dir": "tasks",
            "use_docker": False,
        },
    )

responses = user.sequential_run([
    {
        "chat_id": 1,
        "recipient": financial_assistant,
        "message": financial_tasks[0],
        "silent": False,
        "summary_method": "reflection_with_llm",
    },
    {
        "chat_id": 2,
        "prerequisites": [1],
        "recipient": research_assistant,
        "message": financial_tasks[1],
        "silent": False,
        "summary_method": "reflection_with_llm",
    },
    {"chat_id": 3, "prerequisites": [1, 2], "recipient": writer, "silent": False, "message": writing_tasks[0]},
])

for response in responses:
    response.process()

for response in responses:
    assert len(response.messages) > 0, "Messages should not be empty"
    assert response.last_speaker in ["Financial_assistant", "Researcher", "writer", "User"], (
        "Last speaker should be one of the agents"
    )
    assert response.summary is not None, "Summary should not be None"

# Chat between two comedian agents

# 1. Import our agent class
from autogen import ConversableAgent
from autogen.io.run_response import AsyncRunResponseProtocol

# 2. Define our LLM configuration for OpenAI's GPT-4o mini,
#    uses the OPENAI_API_KEY environment variable
# llm_config = LLMConfig(api_type="openai", model="gpt-4o-mini")

# 3. Create our agents who will tell each other jokes,
#    with Jack ending the chat when Emma says FINISH
with llm_config:
    jack = ConversableAgent(
        "Jack",
        system_message=("Your name is Jack and you are a comedian in a two-person comedy show."),
        is_termination_msg=lambda x: "FINISH" in x["content"],
        human_input_mode="NEVER",
    )
    emma = ConversableAgent(
        "Emma",
        system_message=(
            "Your name is Emma and you are a comedian "
            "in a two-person comedy show. Say the word FINISH "
            "ONLY AFTER you've heard 2 of Jack's jokes."
        ),
        human_input_mode="NEVER",
    )

# 4. Run the chat
response: AsyncRunResponseProtocol = await jack.a_run(
    emma, message="Emma, tell me a joke about goldfish and peanut butter.", summary_method="reflection_with_llm"
)

async for event in response.events:
    print(event)

    if event.type == "input_request":
        await event.content.respond(input())

print(f"{await response.summary=}")
print(f"{await response.messages=}")

# Chat between two comedian agents

# 1. Import our agent class
from autogen import ConversableAgent
from autogen.io.run_response import AsyncRunResponseProtocol

# 2. Define our LLM configuration for OpenAI's GPT-4o mini,
#    uses the OPENAI_API_KEY environment variable
# llm_config = LLMConfig(api_type="openai", model="gpt-4o-mini")

# 3. Create our agents who will tell each other jokes,
#    with Jack ending the chat when Emma says FINISH
with llm_config:
    jack = ConversableAgent(
        "Jack",
        system_message=("Your name is Jack and you are a comedian in a two-person comedy show."),
        is_termination_msg=lambda x: "FINISH" in x["content"],
        human_input_mode="NEVER",
    )
    emma = ConversableAgent(
        "Emma",
        system_message=(
            "Your name is Emma and you are a comedian "
            "in a two-person comedy show. Say the word FINISH "
            "ONLY AFTER you've heard 2 of Jack's jokes."
        ),
        human_input_mode="NEVER",
    )

# 4. Run the chat
response: AsyncRunResponseProtocol = await jack.a_run(
    emma, message="Emma, tell me a joke about goldfish and peanut butter.", summary_method="reflection_with_llm"
)

await response.process()

print(f"{await response.summary=}")
print(f"{await response.messages=}")

assert await response.last_speaker in ["Jack", "Emma"], "Last speaker should be one of the agents"
assert await response.summary is not None, "Summary should not be None"
assert len(await response.messages) > 0, "Messages should not be empty"

# Group chat amongst agents to create a 4th grade lesson plan
# Flow determined by Group Chat Manager automatically, and
# should be Teacher > Planner > Reviewer > Teacher (repeats if necessary)

# 1. Import our agent and group chat classes
from autogen import GroupChat, GroupChatManager

with llm_config:
    # Planner agent setup
    planner_message = "Create lesson plans for 4th grade. Use format: <title>, <learning_objectives>, <script>"
    planner = ConversableAgent(name="planner_agent", system_message=planner_message, description="Creates lesson plans")

    # Reviewer agent setup
    reviewer_message = "Review lesson plans against 4th grade curriculum. Provide max 3 changes."
    reviewer = ConversableAgent(
        name="reviewer_agent", system_message=reviewer_message, description="Reviews lesson plans"
    )

    # Teacher agent setup
    teacher_message = "Choose topics and work with planner and reviewer. Say DONE! when finished."
    teacher = ConversableAgent(
        name="teacher_agent",
        system_message=teacher_message,
    )

# Setup group chat
groupchat = GroupChat(agents=[teacher, planner, reviewer], speaker_selection_method="auto", messages=[])

# Create manager
# At each turn, the manager will check if the message contains DONE! and end the chat if so
# Otherwise, it will select the next appropriate agent using its LLM
manager = GroupChatManager(
    name="group_manager",
    groupchat=groupchat,
    llm_config=llm_config,
    is_termination_msg=lambda x: "DONE!" in (x.get("content", "") or "").upper(),
)

# Start the conversation
response = await teacher.a_run(
    recipient=manager, message="Let's teach the kids about the solar system.", summary_method="reflection_with_llm"
)

await response.process()

assert await response.summary is not None, "Summary should not be None"
assert len(await response.messages) > 0, "Messages should not be empty"
assert await response.last_speaker in ["teacher_agent", "planner_agent", "reviewer_agent"], (
    "Last speaker should be one of the agents"
)

from autogen import a_run_swarm
from autogen.agentchat.contrib.swarm_agent import AfterWorkOption

# 1. Create our agents
planner_message = """You are a classroom lesson planner.
Given a topic, write a lesson plan for a fourth grade class.
If you are given revision feedback, update your lesson plan and record it.
Use the following format:
<title>Lesson plan title</title>
<learning_objectives>Key learning objectives</learning_objectives>
<script>How to introduce the topic to the kids</script>
"""

reviewer_message = """You are a classroom lesson reviewer.
You compare the lesson plan to the fourth grade curriculum
and provide a maximum of 3 recommended changes for each review.
Make sure you provide recommendations each time the plan is updated.
"""

teacher_message = """You are a classroom teacher.
You decide topics for lessons and work with a lesson planner.
and reviewer to create and finalise lesson plans.
"""

with llm_config:
    lesson_planner = ConversableAgent(name="planner_agent", system_message=planner_message)

    lesson_reviewer = ConversableAgent(name="reviewer_agent", system_message=reviewer_message)

    teacher = ConversableAgent(
        name="teacher_agent",
        system_message=teacher_message,
    )

# 2. Initiate the swarm chat using a swarm manager who will
# select agents automatically
response = await a_run_swarm(
    initial_agent=teacher,
    agents=[lesson_planner, lesson_reviewer, teacher],
    messages="Today, let's introduce our kids to the solar system.",
    max_rounds=10,
    swarm_manager_args={"llm_config": llm_config},
    after_work=AfterWorkOption.SWARM_MANAGER,
)

# await response.process()

async for event in response.events:
    if event.type == "input_request":
        await event.content.respond("exit")

assert await response.summary is not None, "Summary should not be None"
assert len(await response.messages) > 0, "Messages should not be empty"
assert await response.last_speaker in ["teacher_agent", "planner_agent", "reviewer_agent"], (
    "Last speaker should be one of the agents"
)

# 1. Create our LLM agent
with llm_config:
    my_agent = ConversableAgent(
        name="helpful_agent",
        system_message="You are a poetic AI assistant, respond in rhyme.",
    )

# 2. Run the agent with a prompt
response = await my_agent.a_run(
    message="In one sentence, what's the big deal about AI?", max_turns=1, summary_method="reflection_with_llm"
)

await response.process()
print(f"{await response.summary=}")
print(f"{await response.messages=}")
print(f"{await response.last_speaker=}")

assert await response.summary is not None, "Summary should not be None"
assert len(await response.messages) == 2, "Messages should not be empty"
assert await response.last_speaker == "helpful_agent", "Last speaker should be an agent"

financial_tasks = [
    """What are the current stock prices of NVDA and TESLA, and how is the performance over the past month in terms of percentage change?""",
    """Investigate possible reasons of the stock performance.""",
]

writing_tasks = ["""Develop an engaging blog post using any information provided."""]

with llm_config:
    financial_assistant = ConversableAgent(
        name="Financial_assistant",
        system_message="You are a financial assistant, helping with stock market analysis. Reply 'TERMINATE' when financial tasks are done.",
        llm_config=llm_config,
    )
    research_assistant = ConversableAgent(
        name="Researcher",
        system_message="You are a research assistant, helping with stock market analysis. Reply 'TERMINATE' when research tasks are done.",
        llm_config=llm_config,
    )
    writer = ConversableAgent(
        name="writer",
        llm_config=llm_config,
        system_message="""
            You are a professional writer, known for
            your insightful and engaging articles.
            You transform complex concepts into compelling narratives.
            Reply "TERMINATE" in the end when everything is done.
            """,
    )

    user = UserProxyAgent(
        name="User",
        human_input_mode="NEVER",
        is_termination_msg=lambda x: x.get("content", "") and x.get("content", "").rstrip().endswith("TERMINATE"),
        code_execution_config={
            "last_n_messages": 1,
            "work_dir": "tasks",
            "use_docker": False,
        },
    )

responses = await user.a_sequential_run([
    {
        "chat_id": 1,
        "recipient": financial_assistant,
        "message": financial_tasks[0],
        "silent": False,
        "summary_method": "reflection_with_llm",
    },
    {
        "chat_id": 2,
        "prerequisites": [1],
        "recipient": research_assistant,
        "message": financial_tasks[1],
        "silent": False,
        "summary_method": "reflection_with_llm",
    },
    {"chat_id": 3, "prerequisites": [1, 2], "recipient": writer, "silent": False, "message": writing_tasks[0]},
])

for response in responses:
    async for event in response.events:
        print(event)
        if event.type == "input_request":
            await event.content.respond("exit")

    assert len(await response.messages) > 0, "Messages should not be empty"
    assert await response.last_speaker in ["Financial_assistant", "Researcher", "writer", "User"], (
        "Last speaker should be one of the agents"
    )
    assert await response.summary is not None, "Summary should not be None"

from autogen.agentchat.group.patterns.auto import AutoPattern

with llm_config:
    triage_agent = ConversableAgent(
        name="triage_agent",
        system_message="""You are a triage agent. For each user query,
        identify whether it is a technical issue or a general question. Route
        technical issues to the tech agent and general questions to the general agent.
        Do not provide suggestions or answers, only route the query.""",
    )

    tech_agent = ConversableAgent(
        name="tech_agent",
        system_message="""You solve technical problems like software bugs
        and hardware issues.""",
    )

    general_agent = ConversableAgent(
        name="general_agent", system_message="You handle general, non-technical support questions."
    )

user = ConversableAgent(name="user", human_input_mode="ALWAYS")

pattern = AutoPattern(
    initial_agent=triage_agent,
    agents=[triage_agent, tech_agent, general_agent],
    user_agent=user,
    group_manager_args={"llm_config": llm_config},
)

# Run the chat

from autogen.agentchat import run_group_chat
from autogen.io.run_response import Cost

response = run_group_chat(
    pattern=pattern, messages="My laptop keeps shutting down randomly. Can you help?", max_rounds=15
)

for event in response.events:
    print(event)
    if event.type == "input_request":
        event.content.respond("exit")

assert response.summary is not None, "Summary should not be None"
assert len(response.messages) > 0, "Messages should not be empty"
assert response.last_speaker in ["tech_agent", "general_agent", "triage_agent"]
assert isinstance(response.cost, Cost)

# Run the chat

from autogen.agentchat import a_run_group_chat

response = await a_run_group_chat(
    pattern=pattern, messages="My laptop keeps shutting down randomly. Can you help?", max_rounds=15
)

async for event in response.events:
    print(event)
    if event.type == "input_request":
        await event.content.respond("exit")

assert await response.summary is not None, "Summary should not be None"
assert len(await response.messages) > 0, "Messages should not be empty"
assert await response.last_speaker in ["tech_agent", "general_agent", "triage_agent"]
assert isinstance(await response.cost, Cost)