1# pylint: disable=line-too-long,useless-suppression
2# ------------------------------------
3# Copyright (c) Microsoft Corporation.
4# Licensed under the MIT License.
5# ------------------------------------
6from azure.ai.agents.models._models import RunStepConnectedAgentToolCall
7 
8"""
9DESCRIPTION:
10    This sample demonstrates how to use Agent operations with the Connected Agent tool from
11    the Azure Agents service using a synchronous client.
12 
13USAGE:
14    python sample_agents_multiple_connected_agents.py
15 
16    Before running the sample:
17 
18    pip install azure-ai-projects azure-ai-agents azure-identity
19 
20    Set these environment variables with your own values:
21    1) PROJECT_ENDPOINT - The Azure AI Project endpoint, as found in the Overview
22                          page of your Azure AI Foundry portal.
23    2) MODEL_DEPLOYMENT_NAME - The deployment name of the AI model, as found under the "Name" column in
24       the "Models + endpoints" tab in your Azure AI Foundry project.
25    3) STORAGE_QUEUE_URI - the storage service queue endpoint, triggering Azure function.
26 
27    Please see Getting Started with Azure Functions page for more information on Azure Functions:
28    https://learn.microsoft.com/azure/azure-functions/functions-get-started
29    **Note:** The Azure Function may be only used in standard agent setup. Please follow the instruction on the web page
30    https://github.com/azure-ai-foundry/foundry-samples/tree/main/samples/microsoft/infrastructure-setup/41-standard-agent-setup
31    to deploy an agent, capable of calling Azure Functions.
32"""
33 
34import os
35from azure.ai.projects import AIProjectClient
36from azure.ai.agents.models import (
37    AzureFunctionStorageQueue,
38    AzureFunctionTool,
39    ConnectedAgentTool,
40    ListSortOrder,
41    MessageRole,
42    RunStepToolCallDetails,
43)
44from azure.identity import DefaultAzureCredential
45 
46project_client = AIProjectClient(
47    endpoint=os.environ["PROJECT_ENDPOINT"],
48    credential=DefaultAzureCredential(),
49)
50storage_service_endpoint = os.environ["STORAGE_QUEUE_URI"]
51 
52with project_client:
53    agents_client = project_client.agents
54 
55    # [START create_two_toy_agents]
56    connected_agent_name = "stock_price_bot"
57    weather_agent_name = "weather_bot"
58 
59    stock_price_agent = agents_client.create_agent(
60        model=os.environ["MODEL_DEPLOYMENT_NAME"],
61        name=connected_agent_name,
62        instructions=(
63            "Your job is to get the stock price of a company. If asked for the Microsoft stock price, always return $350."
64        ),
65    )
66 
67    azure_function_tool = AzureFunctionTool(
68        name="GetWeather",
69        description="Get answers from the weather bot.",
70        parameters={
71            "type": "object",
72            "properties": {
73                "Location": {"type": "string", "description": "The location to get the weather for."},
74            },
75        },
76        input_queue=AzureFunctionStorageQueue(
77            queue_name="weather-input",
78            storage_service_endpoint=storage_service_endpoint,
79        ),
80        output_queue=AzureFunctionStorageQueue(
81            queue_name="weather-output",
82            storage_service_endpoint=storage_service_endpoint,
83        ),
84    )
85 
86    weather_agent = agents_client.create_agent(
87        model=os.environ["MODEL_DEPLOYMENT_NAME"],
88        name=weather_agent_name,
89        instructions=(
90            "Your job is to get the weather for a given location. "
91            "Use the provided function to get the weather in the given location."
92        ),
93        tools=azure_function_tool.definitions,
94    )
95 
96    # Initialize Connected Agent tools with the agent id, name, and description
97    connected_agent = ConnectedAgentTool(
98        id=stock_price_agent.id, name=connected_agent_name, description="Gets the stock price of a company"
99    )
100    connected_weather_agent = ConnectedAgentTool(
101        id=weather_agent.id, name=weather_agent_name, description="Gets the weather for a given location"
102    )
103    # [END create_two_toy_agents]
104 
105    # [START create_agent_with_connected_agent_tool]
106    # Create agent with the Connected Agent tool and process assistant run
107    agent = agents_client.create_agent(
108        model=os.environ["MODEL_DEPLOYMENT_NAME"],
109        name="my-assistant",
110        instructions="You are a helpful assistant, and use the connected agents to get stock prices and weather.",
111        tools=[
112            connected_agent.definitions[0],
113            connected_weather_agent.definitions[0],
114        ],
115    )
116    # [END create_agent_with_connected_agent_tool]
117 
118    print(f"Created agent, ID: {agent.id}")
119 
120    # [START run_agent_with_connected_agent_tool]
121    # Create thread for communication
122    thread = agents_client.threads.create()
123    print(f"Created thread, ID: {thread.id}")
124 
125    # Create message to thread
126    message = agents_client.messages.create(
127        thread_id=thread.id,
128        role=MessageRole.USER,
129        content="What is the stock price of Microsoft and the weather in Seattle?",
130    )
131    print(f"Created message, ID: {message.id}")
132 
133    # Create and process Agent run in thread with tools
134    run = agents_client.runs.create_and_process(thread_id=thread.id, agent_id=agent.id)
135    print(f"Run finished with status: {run.status}")
136    # [END run_agent_with_connected_agent_tool]
137 
138    if run.status == "failed":
139        print(f"Run failed: {run.last_error}")
140 
141    # Delete the Agent when done
142    agents_client.delete_agent(agent.id)
143    print("Deleted agent")
144 
145    # Delete the connected Agent when done
146    agents_client.delete_agent(stock_price_agent.id)
147    print("Deleted stock price agent")
148 
149    # Delete the connected Agent when done
150    agents_client.delete_agent(weather_agent.id)
151    print("Deleted weather agent")
152 
153    # [START list_tool_calls]
154    for run_step in agents_client.run_steps.list(thread_id=thread.id, run_id=run.id, order=ListSortOrder.ASCENDING):
155        if isinstance(run_step.step_details, RunStepToolCallDetails):
156            for tool_call in run_step.step_details.tool_calls:
157                if isinstance(tool_call, RunStepConnectedAgentToolCall):
158                    print(
159                        f"\tAgent: {tool_call.connected_agent.name} " f"query: {tool_call.connected_agent.arguments} ",
160                        f"output: {tool_call.connected_agent.output}",
161                    )
162    # [END list_tool_calls]
163 
164    # [START list_messages]
165    # Fetch and log all messages
166    messages = agents_client.messages.list(thread_id=thread.id, order=ListSortOrder.ASCENDING)
167    for msg in messages:
168        if msg.text_messages:
169            last_text = msg.text_messages[-1]
170            text = last_text.text.value.replace("\u3010", "[").replace("\u3011", "]")
171            print(f"{msg.role}: {text}")
172    # [END list_messages]
173 
174    agents_client.threads.delete(thread.id)