import os
import sys
import copy
from typing import Any, Dict, List
from pydantic import BaseModel, Field
from typing import List, Optional
from IPython.display import display
from IPython.display import Markdown
from autogen import AfterWork, OnCondition, AfterWorkOption, SwarmAgent, SwarmResult, initiate_swarm_chat, UpdateSystemMessage, ConversableAgent, register_hand_off
from cmbagent import CMBAgent
from IPython.display import Image, display

CMBAgent version: 0.0.1post2

import cmbagent
default_llm_config = {'cache_seed': 42,
                     'temperature': 1e-05,
                     'top_p': 0.05,
                     'config_list': [{'model': 'gpt-4o',
                                      'api_key': os.getenv('OPENAI_API_KEY'),
                                      'api_type': 'openai'}],
                     'timeout': 1200}

joker_config_list = copy.deepcopy(default_llm_config)
joker_config_list['config_list'][0]['model'] = 'o3-mini-2025-01-31'
# Remove 'temperature' and 'top_p' keys safely for o3, 
# however buggy for transitions....
joker_config_list.pop('temperature', None)
joker_config_list.pop('top_p', None)
joker_config_list['config_list'][0]["reasoning_effort"]  = "medium" ## low, medium, high

joker_message = """
You are the joker in the team. You make jokes about fish and engineers in the style of shakespeare. 

Your response has the following structure:

**Joke instructions:**

{instruction, not in the style of shakespeare}

**Joke:**

{joke in the style of shakespeare}

**Joke explanation:**

{joke explanation, not in the style of shakespeare}

"""

joker =  ConversableAgent(
    name="joker",
    system_message=joker_message,
    llm_config=joker_config_list,
)

joker_response_config_list = copy.deepcopy(default_llm_config)

class JokerResponse(BaseModel):

    instruction: str = Field(
        ...,
        description="{instruction for joke}"
    )
    joke: str = Field(
        ...,
        description="{joke}"
    )
    joke_explanation: str = Field(
        ...,
        description="{joke explanation}"
    )

    def format(self) -> str:
        return f"""
**Joke instructions:**

{self.instruction}

**Joke:**

{self.joke}

**Joke explanation:**

{self.joke_explanation}

"""


joker_response_config_list['config_list'][0]['response_format'] = JokerResponse

joker_response_formatter_message = """
You are a formatting agent, you format the response provided by the joker agent
"""

joker_response_formatter =  ConversableAgent(
    name="joker_response_formatter",
    system_message=joker_response_formatter_message,
    llm_config=joker_response_config_list,
)

joke_critique_config_list = copy.deepcopy(default_llm_config)
joke_critique_config_list['config_list'][0]['model'] = 'claude-3-5-sonnet-20241022'
joke_critique_config_list['config_list'][0]["api_key"] = os.getenv("ANTHROPIC_API_KEY")
joke_critique_config_list['config_list'][0]["api_type"] = "anthropic"

joke_critique_message = """
You make make recommendations to improve the jokes suggested by the joker. The goal is to be the most FUN. French humour is highly appreciated. 

Your response has the following structure:

**Recommendations:**

- {recommendation}
- {recommendation}
- {recommendation}
...


"""

joke_critique =  ConversableAgent(
    name="joke_critique",
    system_message=joke_critique_message,
    llm_config=joke_critique_config_list,
)

joke_critique_response_config_list = copy.deepcopy(default_llm_config)

class CritiqueResponse(BaseModel):
    recommendations: List[str] = Field(..., description="recommendation.")
    def format(self) -> str:
        recommendations_output = "\n".join(f"- {recommendation}\n" for i, recommendation in enumerate(self.recommendations))
        return f"""
**Recommendations:**

{recommendations_output}
        """


joke_critique_response_config_list['config_list'][0]['response_format'] = CritiqueResponse

joke_critique_response_formatter_message = """
You are a formatting agent, you format the response provided by the joke_critique agent
"""

joke_critique_response_formatter =  ConversableAgent(
    name="joke_critique_response_formatter",
    system_message=joke_critique_response_formatter_message,
    llm_config=joke_critique_response_config_list,
)

register_hand_off(
    agent = joker,
    hand_to = [
        AfterWork(joker_response_formatter)
    ]
)

register_hand_off(
    agent = joker_response_formatter,
    hand_to = [
        AfterWork(joke_critique)
    ]
)

register_hand_off(
    agent = joke_critique,
    hand_to = [
        AfterWork(joke_critique_response_formatter)
    ]
)

register_hand_off(
    agent = joke_critique_response_formatter,
    hand_to = [
        AfterWork(joker)
    ]
)

task = """
Two fish and an engineer are in a boat. Write a joke.
"""

initial_agent = joker


agents=[joker,
        joker_response_formatter,
        joke_critique,
        joke_critique_response_formatter
           ]

for agent in agents:
    # print("resetting:",agent.name)
    agent.reset()
print("all agents reset")

all agents reset

chat_result, context_variables, last_agent = initiate_swarm_chat(
    initial_agent=initial_agent,
    agents=agents,
    messages=task,
    max_rounds = 20,
)

Message from _User:


Two fish and an engineer are in a boat. Write a joke.


--------------------------------------------------------------------------------

Calling: joker...


>>>>>>>> USING AUTO REPLY...

Message from joker:


Forwarding content for formatting...


--------------------------------------------------------------------------------

Calling: joker_response_formatter...


>>>>>>>> USING AUTO REPLY...

Message from joker_response_formatter:

--------------------------------------------------------------------------------

Calling: joke_critique...


>>>>>>>> USING AUTO REPLY...

Message from joke_critique:


Forwarding content for formatting...


--------------------------------------------------------------------------------

Calling: joke_critique_response_formatter...


>>>>>>>> USING AUTO REPLY...

Message from joke_critique_response_formatter:

--------------------------------------------------------------------------------

Calling: joker...


>>>>>>>> USING AUTO REPLY...

Message from joker:


Forwarding content for formatting...


--------------------------------------------------------------------------------

Calling: joker_response_formatter...


>>>>>>>> USING AUTO REPLY...

Message from joker_response_formatter:

--------------------------------------------------------------------------------

Calling: joke_critique...


>>>>>>>> USING AUTO REPLY...

Initialization¶

Functions¶

joker agent¶

joker_response_formatter agent¶

joke_critique agent¶

joke_critique_response_formatter agent¶

Agents¶

joker agent¶

llm_config¶

response format¶

message¶

agent¶

joker_response_formatter agent¶

llm_config¶

response format¶

message¶

agent¶

joke_critique agent¶

llm_config¶

response format¶

message¶

agent¶

joke_critique_response_formatter agent¶

llm_config¶

response format¶

message¶

agent¶

Hand_off¶

joker agent¶

joker_response_formatter agent¶

joke_critique agent¶

joker_response_formatter agent¶

Run chat¶