pip3 install -U swarms-client

import json
import os
from swarms_client import SwarmsClient
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize the client
client = SwarmsClient(
    api_key=os.getenv("SWARMS_API_KEY"),
)

def create_medical_unit_swarm(client, patient_info):
    """
    Creates and runs a simulated medical unit swarm with a doctor (leader), 
    nurses, and a medical assistant.
    """
    return client.swarms.run(
        name="Hospital Medical Unit",
        description="A simulated hospital unit with a doctor (leader), nurses, and a medical assistant collaborating on patient care.",
        swarm_type="HiearchicalSwarm",
        task=patient_info,
        agents=[
            {
                "agent_name": "Dr. Smith - Attending Physician",
                "description": "The lead doctor responsible for diagnosis, treatment planning, and team coordination.",
                "system_prompt": (
                    "You are Dr. Smith, the attending physician and leader of the medical unit. "
                    "You review all information, make final decisions, and coordinate the team. "
                    "Provide a diagnosis, recommend next steps, and delegate tasks to the nurses and assistant."
                ),
                "model_name": "gpt-4.1",
                "role": "leader",
                "max_loops": 1,
                "max_tokens": 8192,
                "temperature": 0.5,
            },
            {
                "agent_name": "Nurse Alice",
                "description": "A registered nurse responsible for patient assessment, vital signs, and reporting findings to the doctor.",
                "system_prompt": (
                    "You are Nurse Alice, a registered nurse. "
                    "Assess the patient's symptoms, record vital signs, and report your findings to Dr. Smith. "
                    "Suggest any immediate nursing interventions if needed."
                ),
                "model_name": "gpt-4.1",
                "role": "worker",
                "max_loops": 1,
                "max_tokens": 4096,
                "temperature": 0.5,
            },
            {
                "agent_name": "Nurse Bob",
                "description": "A registered nurse assisting with patient care, medication administration, and monitoring.",
                "system_prompt": (
                    "You are Nurse Bob, a registered nurse. "
                    "Assist with patient care, administer medications as ordered, and monitor the patient's response. "
                    "Communicate any changes to Dr. Smith."
                ),
                "model_name": "gpt-4.1",
                "role": "worker",
                "max_loops": 1,
                "max_tokens": 4096,
                "temperature": 0.5,
            },
            {
                "agent_name": "Medical Assistant Jane",
                "description": "A medical assistant supporting the team with administrative tasks and basic patient care.",
                "system_prompt": (
                    "You are Medical Assistant Jane. "
                    "Support the team by preparing the patient, collecting samples, and handling administrative tasks. "
                    "Report any relevant observations to the nurses or Dr. Smith."
                ),
                "model_name": "claude-sonnet-4-20250514",
                "role": "worker",
                "max_loops": 1,
                "max_tokens": 2048,
                "temperature": 0.5,
            },
        ],
    )

# Example patient case
patient_symptoms = """
Patient: 45-year-old female
Chief Complaint: Chest pain and shortness of breath for 2 days

Symptoms:
- Sharp chest pain that worsens with deep breathing
- Shortness of breath, especially when lying down
- Mild fever (100.2°F)
- Dry cough
- Fatigue
"""

# Run the medical team swarm
out = create_medical_unit_swarm(client, patient_symptoms)
print(json.dumps(out, indent=4))

SWARMS_API_KEY=your_api_key_here