# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------

"""
DESCRIPTION:
    This sample demonstrates how to use multiple agents using AgentTeam with traces.

    The team consists of
        • one leader agent - automatically created by AgentTeam from the
          configuration in `utils/agent_team_config.yaml`
        • two worker agents - `Coder` and `Reviewer`, defined in the code below

    IMPORTANT - leader-agent model configuration
        `utils/agent_team_config.yaml` contains the key  TEAM_LEADER_MODEL.
        Its value must be the name of a **deployed** model in your Azure AI
        project (e.g. "gpt-4o-mini").
        If this model deployment is not available, AgentTeam cannot instantiate
        the leader agent and the sample will fail.

USAGE:
    python sample_agents_agent_team.py

    Before running the sample:

    1. pip install azure-ai-projects azure-ai-agents azure-identity
    2. Ensure `utils/agent_team_config.yaml` is present and TEAM_LEADER_MODEL points
       to a valid model deployment.
    3. Set these environment variables with your own values:
         PROJECT_ENDPOINT - The Azure AI Project endpoint, as found in the Overview
                            page of your Azure AI Foundry portal.
         MODEL_DEPLOYMENT_NAME - The model deployment name used for the worker agents.
"""

import os
from azure.ai.projects import AIProjectClient
from azure.identity import DefaultAzureCredential
from utils.agent_team import AgentTeam, _create_task
from utils.agent_trace_configurator import AgentTraceConfigurator

project_client = AIProjectClient(
    endpoint=os.environ["PROJECT_ENDPOINT"],
    credential=DefaultAzureCredential(),
)

with project_client:
    agents_client = project_client.agents

    agents_client.enable_auto_function_calls({_create_task})

    model_deployment_name = os.getenv("MODEL_DEPLOYMENT_NAME")

    if model_deployment_name is not None:
        AgentTraceConfigurator(agents_client=agents_client).setup_tracing()
        with agents_client:
            agent_team = AgentTeam("test_team", agents_client=agents_client)
            agent_team.add_agent(
                model=model_deployment_name,
                name="Coder",
                instructions="You are software engineer who writes great code. Your name is Coder.",
            )
            agent_team.add_agent(
                model=model_deployment_name,
                name="Reviewer",
                instructions="You are software engineer who reviews code. Your name is Reviewer.",
            )
            agent_team.assemble_team()

            print("A team of agents specialized in software engineering is available for requests.")
            while True:
                user_input = input("Input (type 'quit' or 'exit' to exit): ")
                if user_input.lower() == "quit":
                    break
                elif user_input.lower() == "exit":
                    break
                agent_team.process_request(request=user_input)

            agent_team.dismantle_team()
    else:
        print("Error: Please define the environment variable MODEL_DEPLOYMENT_NAME.")
