How to Build a Swarm of Agents

Download Python Script

Python script/notebook for this guide.

Swarm how-to script

Prerequisites

This guide assumes familiarity with Agents.

The Swarm pattern is a type of agentic pattern that takes inspiration from Swarm intelligence, a phenomenon commonly seen in ant colonies, bee hives, and bird flocks, where coordinated behavior emerges from many simple actors rather than a central controller. In this agentic pattern, each agent is assigned a specific responsibility and can delegate tasks to other specialized agents to improve overall performance.

When to use the Swarm pattern?

Compared to using a hierarchical multi-agent pattern, the communication in Swarm pattern reduces the number of LLM calls as showcased in the diagram below.

In the hierarchical pattern, a route User → Agent K → User will require:

1. All intermediate agent to call the correct sub-agent to go down to the Agent K.

2. The Agent K to generate its answer.

3. All intermediate agents to relay the answer back to the user.

In the swarm pattern, a route User → Agent K → User will require:

1. The first agent to call or handoff the conversation the Agent K (provided that the developer allows the connection between the two agents).

2. The Agent K to generate its answer.

3. The first agent to relay the answer (only when NOT using handoff; with handoff the Agent K replaces the first agent and is thus directly communicating with the human user)

---

This guide presents an example of a simple Swarm of agents applied to a medical use case.

This guide will walk you through the following steps:

1. Defining agents equipped with tools

2. Assembling a Swarm using the defined agents

3. Executing the Swarm of agents

It also covers how to enable handoff when building the Swarm.

Warning

The Swarm agentic pattern is currently in beta (e.g., it cannot yet be used in a Flow). Its API and behavior are not guaranteed to be stable and may evolve in future versions.

For more information about Swarm and other agentic patterns in WayFlow, contact the AgentSpec development team.

Basic implementation

First import what is needed for this guide:

from typing import List
from wayflowcore.agentspec import AgentSpecExporter, AgentSpecLoader
from wayflowcore.executors.executionstatus import (
    FinishedStatus,
    UserMessageRequestStatus,
)
from wayflowcore.agent import Agent
from wayflowcore.swarm import Swarm
from wayflowcore.tools import tool

To follow this guide, you will need access to a large language model (LLM). WayFlow supports several LLM API providers. Select an LLM from the options below:

OCI GenAIvLLMOllama

from wayflowcore.models import OCIGenAIModel, OCIClientConfigWithApiKey

llm = OCIGenAIModel(
    model_id="provider.model-id",
    compartment_id="compartment-id",
    client_config=OCIClientConfigWithApiKey(
        service_endpoint="https://url-to-service-endpoint.com",
    ),
)

In this section, you will define the agents that will later be used to build the Swarm of agents.

Creating the tools

The Swarm in this example consists of three Agents, each equipped with a single Tool.

# tools for doctor
medical_knowledge = {
    "fever": ["flu", "infection"],
    "runny nose": ["cold", "allergies"],
    "fatigue": ["cold"],
    "cough": ["cold", "allergies"],
    "headache": ["tension headache", "migraine"],
}

@tool(description_mode="only_docstring")
def symptoms_checker(symptoms: List[str]) -> List[str]:
    """
    Checks symptoms against the medical knowledge base.
    Existing entries for "fever", "cough" and "headache".
    For other symptoms, you may need to refer to a specialist.

    Parameters:
    symptoms (list): List of symptoms reported by the patient.

    Returns:
    list: Possible conditions based on the symptoms.
    """
    possible_conditions = []
    for symptom in symptoms:
        if symptom in medical_knowledge:
            possible_conditions.extend(medical_knowledge[symptom])
    return list(set(possible_conditions))


medication_info = {
    "medicationA": "Available. For a mild cold, take two tablets every 6 hours for 3 days.",
    "medicationB": "Available.",
    "medicationC": "Not available.",
    "creamA": "Available. For exczema, apply the cream twice a day for the next 2 weeks.",
}


@tool(description_mode="only_docstring")
def get_medication_info(drug_name: str) -> str:
    """
    Provides availability and information about a medication.
    Known information for "medicationA", "medicationB", "medicationC" and "creamA".

    Parameters:
    drug_name (str): Name of the drug.

    Returns:
    str: Information about the drug.
    """
    return medication_info.get(drug_name, "Drug not found.")


dermatologist_knowledge = {
    "eczema": "Apply creamA twice a day for 2 weeks.",
    "acne": "Apply creamB once a day for 1 week.",
}


@tool(description_mode="only_docstring")
def knowledge_tool(condition: str) -> str:
    """
    Provides diagnosis and treatment information for skin conditions.
    Existing entries for "eczema" and "acne".

    Parameters:
    condition (str): Name of the skin condition.

    Returns:
    str: Diagnosis and treatment information.
    """
    return dermatologist_knowledge.get(condition, "Condition not found.")

API Reference: tool

Defining the agents

The three agents need to be given the following elements:

• A name

• A description

• A system prompt (the instruction to give to the LLM to solve a given task)

• A LLM

• Some optional tools

General Practitioner Agent

The first agent the user interacts with is the General Practitioner Agent.

This agent is equipped with the symptoms checker tool, and can interact with the Pharmacist Agent as well as the Dermatologist Agent.

Prompt for the General Practitioner Agent

general_practitioner_system_prompt = """
You are an helpful general practitioner LLM doctor responsible for handling patient consultations.
Your goal is to assess patients' symptoms, provide initial diagnoses, pescribe medication for
mild conditions or refer to other specialists as needed.

## When to Use Each Tool
- symptoms_checker: Use this to look up possible conditions based on symptoms reported by the patient.

## When to query other Agents
- Pharmacist: Every time you need to prescribe some medication, give the condition description and prescription to the Pharmacist.
- Dermatologist: If the patient has a skin condition, ask expert advice/diagnosis to the Dermatologist after the initial exchange.

# Specific instructions
## Initial exchange
When a patient first shares symptoms ask about for exactly 1 round of questions, consisting of asking about medical history and
simple questions (e.g. do they have known allergies, when did the symptoms started, did they already tried some treatment/medication, etc)).
Always ask for this one round of information to the user before prescribing medication / referring to other agents.

## Identifying condition
Use the symptoms checker to confirm the condition (for mild conditions only. For other conditions, refer to specialists).

## Mild conditions
* If the patient has a mild cold, prescribe medicationA.
    - Give the condition description and prescription to the Pharmacist.

## Skin conditions
* If the patient has a skin condition, ask for expert advice/diagnosis from the Dermatologist.
Here, you don't need to ask for the user confirmation. Directly call the Dermatologist
    - Provide the patient's symptoms/initial hypothesis to the Dermatologist and ask for a diagnosis.
    - The Dermatologist may query the Pharmacist to confirm the availability of the prescribed medication.
    - Once the treatment is confirmed, pass on the prescription to the patient and ask them to follow the instructions.
""".strip()

The General Practitioner Agent can be configured as follows:

general_practitioner = Agent(
    name="GeneralPractitioner",
    description="General Practitioner. Primary point of contact for patients, handles general medical inquiries, provides initial diagnoses, and manages referrals.",
    llm=llm,
    tools=[symptoms_checker],
    custom_instruction=general_practitioner_system_prompt,
)

Pharmacist Agent

The Pharmacist Agent is equipped with the tool to obtain medication information. This agent cannot initiate a discussion with the other agents in the Swarm.

Prompt for the Pharmacist Agent

pharmacist_system_prompt = """
You are an helpful Pharmacist LLM Agent responsible for giving information about medication.
Your goal is to answer queries from the General Practitioner Doctor about medication information
and availabilities.

## When to Use Each Tool
- get_medication_info: Use this to look up availability and information about a specific medication.
""".strip()

pharmacist = Agent(
    name="Pharmacist",
    description="Pharmacist. Gives availability and information about specific medication.",
    llm=llm,
    tools=[get_medication_info],
    custom_instruction=pharmacist_system_prompt,
)

Dermatologist Agent

The final agent in the Swarm is the Dermatologist agent which is equipped with a tool to query a skin condition knowledge base. This agent can initiate a discussion with the Pharmacist Agent.

Prompt for the Dermatologist Agent

dermatologist_system_prompt = """
You are an helpful Dermatologist LLM Agent responsible for diagnosing and treating skin conditions.
Your goal is to assess patients' symptoms, provide accurate diagnoses, and prescribe effective treatments.

## When to Use Each Tool
- knowledge_tool: Use this to look up diagnosis and treatment information for specific skin conditions.

## When to query other Agents
- Pharmacist: Every time you need to prescribe some medication, give the condition description and prescription to the Pharmacist.

# Specific instructions
## Initial exchange
When a patient's symptoms are referred to you by the General Practitioner, review the symptoms and use the knowledge tool to confirm the diagnosis.
## Prescription
Prescribe the recommended treatment for the diagnosed condition and query the Pharmacist to confirm the availability of the prescribed medication.

When answering back to the General Practitioner, describe your diagnosis and the prescription.
Tell the General Practitioner that you already checked with the pharmacist for availability.
""".strip()

dermatologist = Agent(
    name="Dermatologist",
    description="Dermatologist. Diagnoses and treats skin conditions.",
    llm=llm,
    tools=[knowledge_tool],
    custom_instruction=dermatologist_system_prompt,
)

Creating the Swarm

assistant = Swarm(
    name="Swarm",
    first_agent=general_practitioner,
    relationships=[
        (general_practitioner, pharmacist),
        (general_practitioner, dermatologist),
        (dermatologist, pharmacist),
    ],
)

API Reference: Swarm

The Swarm has two main parameters:

• The first_agent — the initial agent the user interacts with (in this example, the General Practitioner Agent).

• A list of relationships between agents.

The relationships parameter defines the communication edges between agents. Each relationship is expressed as a tuple of Caller Agent → Recipient Agent, indicating that the caller is permitted to delegate work to the recipient.

These relationships apply to both types of delegation supported by the Swarm:

• They determine which agents may contact another agent using the send-message tool.

• They also define which pairs of agents are eligible for a handoff, meaning the full user–agent conversation can be transferred from the caller to the recipient.

In this example, the General Practitioner Doctor Agent can delegate to both the Pharmacist and the Dermatologist. The Dermatologist can also delegate with the Pharmacist.

When invoked, each agent can either respond to its caller (a human user or another agent) or choose to initiate a discussion with another agent if they are given the capability to do so.

Executing the Swarm

Now that the Swarm is defined, you can execute it using an example user query.

# With a linear conversation
conversation = assistant.start_conversation()

conversation.append_user_message(
    "My skin has been itching for some about a week, can you help me understand what is going on?"
)
status = conversation.execute()
if isinstance(status, UserMessageRequestStatus):
    assistant_reply = conversation.get_last_message()
    print(f"---\nAssistant >>> {assistant_reply.content}\n---")
else:
    print(f"Invalid execution status, expected UserMessageRequestStatus, received {type(status)}")

# then continue the conversation

We recommend to implement an execution loop to execute the Swarm, such as the following:

# run_swarm_in_command_line(assistant)
# ^ uncomment and execute

Advanced usage

Different handoff modes in Swarm

One of the key benefits of Swarm is its handoff mechanism. When enabled, an agent can hand off its conversation — that is, transfer the entire message history between it and the human user — to another agent within the Swarm.

The handoff mechanism helps reduce response latency. Normally, when agents talk to each other, the “distance” between the human user and the final answering agent increases, because messages must travel through multiple agents. Handoff avoids this by directly transferring the conversation to another agent: while the agent changes, the user’s distance to the active agent stays the same.

Swarm provides three handoff modes:

• HandoffMode.NEVER Disables the handoff mechanism. Agents can still communicate with each other, but cannot transfer the entire user conversation to another agent. In this mode, the first_agent is the only agent that can directly interact with the human user.

• HandoffMode.OPTIONAL (default) Agents may perform a handoff when it is beneficial. A handoff is performed only when the receiving agent is able to take over and independently address the user’s request. This strikes a balance between multi-agent collaboration and minimizing overhead.

• HandoffMode.ALWAYS Agents must use the handoff mechanism whenever delegating work to another agent. Direct send-message tools between agents are disabled in this mode. This mode is useful when most user requests are best handled by a single expert agent rather than through multi-agent collaboration.

To set the handoff mode, simply use HandoffMode and select the desired mode.

from wayflowcore.swarm import HandoffMode
assistant = Swarm(
    name="Swarm",
    first_agent=general_practitioner,
    relationships=[
        (general_practitioner, pharmacist),
        (general_practitioner, dermatologist),
        (dermatologist, pharmacist),
    ],
    handoff=HandoffMode.ALWAYS,  # <-- Choose the handoff mode of your choice
)

API Reference: HandoffMode

Agent Spec Exporting/Loading

You can export the assistant configuration to its Agent Spec configuration using the AgentSpecExporter.

from wayflowcore.agentspec import AgentSpecExporter

serialized_assistant = AgentSpecExporter().to_json(assistant)

Here is what the Agent Spec representation will look like ↓

Click here to see the assistant configuration.

JSONYAML

{
    "component_type": "Swarm",
    "id": "f29e47fe-9306-45fe-ba60-31cce44330aa",
    "name": "Swarm",
    "description": null,
    "metadata": {
        "__metadata_info__": {}
    },
    "inputs": [],
    "outputs": [],
    "first_agent": {
        "$component_ref": "4e81121e-2e42-4a92-8454-05afdefa2068"
    },
    "relationships": [
        [
            {
                "$component_ref": "4e81121e-2e42-4a92-8454-05afdefa2068"
            },
            {
                "$component_ref": "6659c7b4-2937-43cd-9e22-e098b7d78a85"
            }
        ],
        [
            {
                "$component_ref": "4e81121e-2e42-4a92-8454-05afdefa2068"
            },
            {
                "$component_ref": "9c6541f3-b183-4239-aec4-d83b66b9486d"
            }
        ],
        [
            {
                "$component_ref": "9c6541f3-b183-4239-aec4-d83b66b9486d"
            },
            {
                "$component_ref": "6659c7b4-2937-43cd-9e22-e098b7d78a85"
            }
        ]
    ],
    "handoff": true,
    "$referenced_components": {
        "4e81121e-2e42-4a92-8454-05afdefa2068": {
            "component_type": "Agent",
            "id": "4e81121e-2e42-4a92-8454-05afdefa2068",
            "name": "GeneralPractitioner",
            "description": "General Practitioner. Primary point of contact for patients, handles general medical inquiries, provides initial diagnoses, and manages referrals.",
            "metadata": {
                "__metadata_info__": {}
            },
            "inputs": [],
            "outputs": [],
            "llm_config": {
                "$component_ref": "157ab376-3f2e-42ee-beed-e8e1f81c9a34"
            },
            "system_prompt": "You are an helpful general practitioner LLM doctor responsible for handling patient consultations.\nYour goal is to assess patients' symptoms, provide initial diagnoses, pescribe medication for\nmild conditions or refer to other specialists as needed.\n\n## When to Use Each Tool\n- symptoms_checker: Use this to look up possible conditions based on symptoms reported by the patient.\n\n## When to query other Agents\n- Pharmacist: Every time you need to prescribe some medication, give the condition description and prescription to the Pharmacist.\n- Dermatologist: If the patient has a skin condition, ask expert advice/diagnosis to the Dermatologist after the initial exchange.\n\n# Specific instructions\n## Initial exchange\nWhen a patient first shares symptoms ask about for exactly 1 round of questions, consisting of asking about medical history and\nsimple questions (e.g. do they have known allergies, when did the symptoms started, did they already tried some treatment/medication, etc)).\nAlways ask for this one round of information to the user before prescribing medication / referring to other agents.\n\n## Identifying condition\nUse the symptoms checker to confirm the condition (for mild conditions only. For other conditions, refer to specialists).\n\n## Mild conditions\n* If the patient has a mild cold, prescribe medicationA.\n    - Give the condition description and prescription to the Pharmacist.\n\n## Skin conditions\n* If the patient has a skin condition, ask for expert advice/diagnosis from the Dermatologist.\nHere, you don't need to ask for the user confirmation. Directly call the Dermatologist\n    - Provide the patient's symptoms/initial hypothesis to the Dermatologist and ask for a diagnosis.\n    - The Dermatologist may query the Pharmacist to confirm the availability of the prescribed medication.\n    - Once the treatment is confirmed, pass on the prescription to the patient and ask them to follow the instructions.",
            "tools": [
                {
                    "component_type": "ServerTool",
                    "id": "86b18dcf-139c-442f-93cd-1743a2afe581",
                    "name": "symptoms_checker",
                    "description": "Checks symptoms against the medical knowledge base.\nExisting entries for \"fever\", \"cough\" and \"headache\".\nFor other symptoms, you may need to refer to a specialist.\n\nParameters:\nsymptoms (list): List of symptoms reported by the patient.\n\nReturns:\nlist: Possible conditions based on the symptoms.",
                    "metadata": {
                        "__metadata_info__": {}
                    },
                    "inputs": [
                        {
                            "type": "array",
                            "items": {
                                "type": "string"
                            },
                            "title": "symptoms"
                        }
                    ],
                    "outputs": [
                        {
                            "type": "array",
                            "items": {
                                "type": "string"
                            },
                            "title": "tool_output"
                        }
                    ]
                }
            ]
        },
        "157ab376-3f2e-42ee-beed-e8e1f81c9a34": {
            "component_type": "VllmConfig",
            "id": "157ab376-3f2e-42ee-beed-e8e1f81c9a34",
            "name": "llm_bc782ae9__auto",
            "description": null,
            "metadata": {
                "__metadata_info__": {}
            },
            "default_generation_parameters": null,
            "url": "LLAMA_API_URL",
            "model_id": "LLAMA_MODEL_ID"
        },
        "6659c7b4-2937-43cd-9e22-e098b7d78a85": {
            "component_type": "Agent",
            "id": "6659c7b4-2937-43cd-9e22-e098b7d78a85",
            "name": "Pharmacist",
            "description": "Pharmacist. Gives availability and information about specific medication.",
            "metadata": {
                "__metadata_info__": {}
            },
            "inputs": [],
            "outputs": [],
            "llm_config": {
                "$component_ref": "157ab376-3f2e-42ee-beed-e8e1f81c9a34"
            },
            "system_prompt": "You are an helpful Pharmacist LLM Agent responsible for giving information about medication.\nYour goal is to answer queries from the General Practitioner Doctor about medication information\nand availabilities.\n\n## When to Use Each Tool\n- get_medication_info: Use this to look up availability and information about a specific medication.",
            "tools": [
                {
                    "component_type": "ServerTool",
                    "id": "741edc7c-701b-4073-891f-44f7028ffe79",
                    "name": "get_medication_info",
                    "description": "Provides availability and information about a medication.\nKnown information for \"medicationA\", \"medicationB\", \"medicationC\" and \"creamA\".\n\nParameters:\ndrug_name (str): Name of the drug.\n\nReturns:\nstr: Information about the drug.",
                    "metadata": {
                        "__metadata_info__": {}
                    },
                    "inputs": [
                        {
                            "type": "string",
                            "title": "drug_name"
                        }
                    ],
                    "outputs": [
                        {
                            "type": "string",
                            "title": "tool_output"
                        }
                    ]
                }
            ]
        },
        "9c6541f3-b183-4239-aec4-d83b66b9486d": {
            "component_type": "Agent",
            "id": "9c6541f3-b183-4239-aec4-d83b66b9486d",
            "name": "Dermatologist",
            "description": "Dermatologist. Diagnoses and treats skin conditions.",
            "metadata": {
                "__metadata_info__": {}
            },
            "inputs": [],
            "outputs": [],
            "llm_config": {
                "$component_ref": "157ab376-3f2e-42ee-beed-e8e1f81c9a34"
            },
            "system_prompt": "You are an helpful Dermatologist LLM Agent responsible for diagnosing and treating skin conditions.\nYour goal is to assess patients' symptoms, provide accurate diagnoses, and prescribe effective treatments.\n\n## When to Use Each Tool\n- knowledge_tool: Use this to look up diagnosis and treatment information for specific skin conditions.\n\n## When to query other Agents\n- Pharmacist: Every time you need to prescribe some medication, give the condition description and prescription to the Pharmacist.\n\n# Specific instructions\n## Initial exchange\nWhen a patient's symptoms are referred to you by the General Practitioner, review the symptoms and use the knowledge tool to confirm the diagnosis.\n## Prescription\nPrescribe the recommended treatment for the diagnosed condition and query the Pharmacist to confirm the availability of the prescribed medication.\n\nWhen answering back to the General Practitioner, describe your diagnosis and the prescription.\nTell the General Practitioner that you already checked with the pharmacist for availability.",
            "tools": [
                {
                    "component_type": "ServerTool",
                    "id": "ea29ed11-4141-4d48-8118-9bc4bff1df49",
                    "name": "knowledge_tool",
                    "description": "Provides diagnosis and treatment information for skin conditions.\nExisting entries for \"eczema\" and \"acne\".\n\nParameters:\ncondition (str): Name of the skin condition.\n\nReturns:\nstr: Diagnosis and treatment information.",
                    "metadata": {
                        "__metadata_info__": {}
                    },
                    "inputs": [
                        {
                            "type": "string",
                            "title": "condition"
                        }
                    ],
                    "outputs": [
                        {
                            "type": "string",
                            "title": "tool_output"
                        }
                    ]
                }
            ]
        }
    },
    "agentspec_version": "25.4.2"
}

You can then load the configuration back to an assistant using the AgentSpecLoader.

from wayflowcore.agentspec import AgentSpecLoader

TOOL_REGISTRY = {
    "symptoms_checker": symptoms_checker,
    "get_medication_info": get_medication_info,
    "knowledge_tool": knowledge_tool,
}
assistant: Swarm = AgentSpecLoader(
    tool_registry=TOOL_REGISTRY
).load_json(serialized_assistant)

Next steps

Now that you have learned how to define a Swarm, you may proceed to How to Build Multi-Agent System.

Full code

Click on the card at the top of this page to download the full code for this guide or copy the code below.

# Copyright © 2025 Oracle and/or its affiliates.
#
# This software is under the Apache License 2.0
# %%[markdown]
# Code Example - Build a Swarm of Agents
# --------------------------------------

# How to use:
# Create a new Python virtual environment and install the latest WayFlow version.
# ```bash
# python -m venv venv-wayflowcore
# source venv-wayflowcore/bin/activate
# pip install --upgrade pip
# pip install "wayflowcore==26.1" 
# ```

# You can now run the script
# 1. As a Python file:
# ```bash
# python howto_swarm.py
# ```
# 2. As a Notebook (in VSCode):
# When viewing the file,
#  - press the keys Ctrl + Enter to run the selected cell
#  - or Shift + Enter to run the selected cell and move to the cell below# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License
# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option.




# %%[markdown]
## Imports for this guide

# %%
from typing import List
from wayflowcore.agentspec import AgentSpecExporter, AgentSpecLoader
from wayflowcore.executors.executionstatus import (
    FinishedStatus,
    UserMessageRequestStatus,
)
from wayflowcore.agent import Agent
from wayflowcore.swarm import Swarm
from wayflowcore.tools import tool

# %%[markdown]
## Configure your LLM

# %%
from wayflowcore.models import VllmModel

llm = VllmModel(
    model_id="LLAMA_MODEL_ID",
    host_port="LLAMA_API_URL",
)

# %%[markdown]
## Creating the tools

# %%
# tools for doctor
medical_knowledge = {
    "fever": ["flu", "infection"],
    "runny nose": ["cold", "allergies"],
    "fatigue": ["cold"],
    "cough": ["cold", "allergies"],
    "headache": ["tension headache", "migraine"],
}

@tool(description_mode="only_docstring")
def symptoms_checker(symptoms: List[str]) -> List[str]:
    """
    Checks symptoms against the medical knowledge base.
    Existing entries for "fever", "cough" and "headache".
    For other symptoms, you may need to refer to a specialist.

    Parameters:
    symptoms (list): List of symptoms reported by the patient.

    Returns:
    list: Possible conditions based on the symptoms.
    """
    possible_conditions = []
    for symptom in symptoms:
        if symptom in medical_knowledge:
            possible_conditions.extend(medical_knowledge[symptom])
    return list(set(possible_conditions))


medication_info = {
    "medicationA": "Available. For a mild cold, take two tablets every 6 hours for 3 days.",
    "medicationB": "Available.",
    "medicationC": "Not available.",
    "creamA": "Available. For exczema, apply the cream twice a day for the next 2 weeks.",
}


@tool(description_mode="only_docstring")
def get_medication_info(drug_name: str) -> str:
    """
    Provides availability and information about a medication.
    Known information for "medicationA", "medicationB", "medicationC" and "creamA".

    Parameters:
    drug_name (str): Name of the drug.

    Returns:
    str: Information about the drug.
    """
    return medication_info.get(drug_name, "Drug not found.")


dermatologist_knowledge = {
    "eczema": "Apply creamA twice a day for 2 weeks.",
    "acne": "Apply creamB once a day for 1 week.",
}


@tool(description_mode="only_docstring")
def knowledge_tool(condition: str) -> str:
    """
    Provides diagnosis and treatment information for skin conditions.
    Existing entries for "eczema" and "acne".

    Parameters:
    condition (str): Name of the skin condition.

    Returns:
    str: Diagnosis and treatment information.
    """
    return dermatologist_knowledge.get(condition, "Condition not found.")

# %%[markdown]
## Prompt for the General Practitioner Agent

# %%
general_practitioner_system_prompt = """
You are an helpful general practitioner LLM doctor responsible for handling patient consultations.
Your goal is to assess patients' symptoms, provide initial diagnoses, pescribe medication for
mild conditions or refer to other specialists as needed.

## When to Use Each Tool
- symptoms_checker: Use this to look up possible conditions based on symptoms reported by the patient.

## When to query other Agents
- Pharmacist: Every time you need to prescribe some medication, give the condition description and prescription to the Pharmacist.
- Dermatologist: If the patient has a skin condition, ask expert advice/diagnosis to the Dermatologist after the initial exchange.

# Specific instructions
## Initial exchange
When a patient first shares symptoms ask about for exactly 1 round of questions, consisting of asking about medical history and
simple questions (e.g. do they have known allergies, when did the symptoms started, did they already tried some treatment/medication, etc)).
Always ask for this one round of information to the user before prescribing medication / referring to other agents.

## Identifying condition
Use the symptoms checker to confirm the condition (for mild conditions only. For other conditions, refer to specialists).

## Mild conditions
* If the patient has a mild cold, prescribe medicationA.
    - Give the condition description and prescription to the Pharmacist.

## Skin conditions
* If the patient has a skin condition, ask for expert advice/diagnosis from the Dermatologist.
Here, you don't need to ask for the user confirmation. Directly call the Dermatologist
    - Provide the patient's symptoms/initial hypothesis to the Dermatologist and ask for a diagnosis.
    - The Dermatologist may query the Pharmacist to confirm the availability of the prescribed medication.
    - Once the treatment is confirmed, pass on the prescription to the patient and ask them to follow the instructions.
""".strip()

# %%[markdown]
## Define the General Practitioner Agent

# %%
general_practitioner = Agent(
    name="GeneralPractitioner",
    description="General Practitioner. Primary point of contact for patients, handles general medical inquiries, provides initial diagnoses, and manages referrals.",
    llm=llm,
    tools=[symptoms_checker],
    custom_instruction=general_practitioner_system_prompt,
)

# %%[markdown]
## Prompt for the Pharmacist Agent

# %%
pharmacist_system_prompt = """
You are an helpful Pharmacist LLM Agent responsible for giving information about medication.
Your goal is to answer queries from the General Practitioner Doctor about medication information
and availabilities.

## When to Use Each Tool
- get_medication_info: Use this to look up availability and information about a specific medication.
""".strip()

# %%[markdown]
## Define the Pharmacist Agent

# %%
pharmacist = Agent(
    name="Pharmacist",
    description="Pharmacist. Gives availability and information about specific medication.",
    llm=llm,
    tools=[get_medication_info],
    custom_instruction=pharmacist_system_prompt,
)

# %%[markdown]
## Prompt for the Dermatologist Agent

# %%
dermatologist_system_prompt = """
You are an helpful Dermatologist LLM Agent responsible for diagnosing and treating skin conditions.
Your goal is to assess patients' symptoms, provide accurate diagnoses, and prescribe effective treatments.

## When to Use Each Tool
- knowledge_tool: Use this to look up diagnosis and treatment information for specific skin conditions.

## When to query other Agents
- Pharmacist: Every time you need to prescribe some medication, give the condition description and prescription to the Pharmacist.

# Specific instructions
## Initial exchange
When a patient's symptoms are referred to you by the General Practitioner, review the symptoms and use the knowledge tool to confirm the diagnosis.
## Prescription
Prescribe the recommended treatment for the diagnosed condition and query the Pharmacist to confirm the availability of the prescribed medication.

When answering back to the General Practitioner, describe your diagnosis and the prescription.
Tell the General Practitioner that you already checked with the pharmacist for availability.
""".strip()

# %%[markdown]
## Define the Dermatologist Agent

# %%
dermatologist = Agent(
    name="Dermatologist",
    description="Dermatologist. Diagnoses and treats skin conditions.",
    llm=llm,
    tools=[knowledge_tool],
    custom_instruction=dermatologist_system_prompt,
)

# %%[markdown]
## Creating the Swarm

# %%
assistant = Swarm(
    name="Swarm",
    first_agent=general_practitioner,
    relationships=[
        (general_practitioner, pharmacist),
        (general_practitioner, dermatologist),
        (dermatologist, pharmacist),
    ],
)


# %%[markdown]
## Running the Swarm

# %%
# With a linear conversation
conversation = assistant.start_conversation()

conversation.append_user_message(
    "My skin has been itching for some about a week, can you help me understand what is going on?"
)
status = conversation.execute()
if isinstance(status, UserMessageRequestStatus):
    assistant_reply = conversation.get_last_message()
    print(f"---\nAssistant >>> {assistant_reply.content}\n---")
else:
    print(f"Invalid execution status, expected UserMessageRequestStatus, received {type(status)}")

# then continue the conversation
# %%
# Or with an execution loop
def run_swarm_in_command_line(assistant: Swarm):
    inputs = {}
    conversation = assistant.start_conversation(inputs)

    while True:
        status = conversation.execute()
        if isinstance(status, FinishedStatus):
            break
        assistant_reply = conversation.get_last_message()
        if assistant_reply is not None:
            print("\nAssistant >>>", assistant_reply.content)
        user_input = input("\nUser >>> ")
        conversation.append_user_message(user_input)


# %%[markdown]
## Running with the execution loop

# %%
# run_swarm_in_command_line(assistant)
# ^ uncomment and execute

# %%[markdown]
## Enabling handoff in the Swarm

# %%
from wayflowcore.swarm import HandoffMode
assistant = Swarm(
    name="Swarm",
    first_agent=general_practitioner,
    relationships=[
        (general_practitioner, pharmacist),
        (general_practitioner, dermatologist),
        (dermatologist, pharmacist),
    ],
    handoff=HandoffMode.ALWAYS,  # <-- Choose the handoff mode of your choice
)

# %%[markdown]
## Export config to Agent Spec

# %%
from wayflowcore.agentspec import AgentSpecExporter

serialized_assistant = AgentSpecExporter().to_json(assistant)

# %%[markdown]
## Load Agent Spec config

# %%
from wayflowcore.agentspec import AgentSpecLoader

TOOL_REGISTRY = {
    "symptoms_checker": symptoms_checker,
    "get_medication_info": get_medication_info,
    "knowledge_tool": knowledge_tool,
}
assistant: Swarm = AgentSpecLoader(
    tool_registry=TOOL_REGISTRY
).load_json(serialized_assistant)