Tutorial: Build a RAG Agent With Azure AI Agent Service SDK
This tutorial will help you build an agent using the Azure AI Python SDK. The agent will complete an action by leveraging the file search tool (knowledge tool) and a user-defined function call (Action tool).
The agent can answer questions about flight timings and baggage policies based on an existing PDF file and a real-time API. While the PDF is fictitious, the API is real-time and based on FlightAware API. Refer to my previous tutorial for details on this API and how to integrate it with function calls.
Prerequisites
This guide assumes that you have access to an active Azure subscription. You will require a hub and project in Azure AI Foundry, with the GPT-4o-mini model deployed and ready to use.
Your Model+Endpoints section should look like the following:
Ensure the token rate limit is set to max to avoid throttling.
Finally, get the project connection string from the project’s resource page.
Install the Python Modules
Configure the virtual environment:
python -m venv .venv source .venv/bin/activate
Install Python modules:
pip install azure-ai-projects pip install azure-identity
Initialize the Client
We will start by initializing the project client and passing the project connection string, which is stored as an environment variable.
import os
from azure.ai.projects import AIProjectClient
from azure.ai.projects.models import FileSearchTool, FilePurpose, FunctionTool, ToolSet
from azure.identity import DefaultAzureCredential
from user_functions import user_functions
# Initialize client
credential = DefaultAzureCredential()
project_client = AIProjectClient.from_connection_string(
credential=credential, conn_str=os.environ["PROJECT_CONNECTION_STRING"]
)
Create the Knowledge Tool From the PDF
We will now upload the PDF to turn it into a vector store, which will become a knowledge tool for the agent.
# Upload and process baggage policy document
file = project_client.agents.upload_file_and_poll(file_path='./data/baggage.pdf', purpose=FilePurpose.AGENTS)
print(f"Uploaded file, file ID: {file.id}")
# Create vector store
vector_store = project_client.agents.create_vector_store_and_poll(file_ids=[file.id], name="baggage_vector_store")
print(f"Created vector store, vector store ID: {vector_store.id}")
Wrapping the FlightAware API Into a Tool
Create a Python file called user_functions.py with the following code:
import json
from datetime import datetime, timedelta
from typing import Any, Callable, Set, Dict, List, Optional
import requests
import pytz
def get_flight_status(flight):
"""Returns Flight Information"""
AEROAPI_BASE_URL = "YOUR_API_URL"
AEROAPI_KEY="YOUR_API_KEY"
def get_api_session():
session = requests.Session()
session.headers.update({"x-apikey": AEROAPI_KEY})
return session
def fetch_flight_data(flight_id, session):
if "flight_id=" in flight_id:
flight_id = flight_id.split("flight_id=")[1]
start_date = datetime.now().date().strftime('%Y-%m-%d')
end_date = (datetime.now().date() + timedelta(days=1)).strftime('%Y-%m-%d')
api_resource = f"/flights/{flight_id}?start={start_date}&end={end_date}"
response = session.get(f"{AEROAPI_BASE_URL}{api_resource}")
response.raise_for_status()
return response.json()['flights'][0]
def utc_to_local(utc_date_str, local_timezone_str):
utc_datetime = datetime.strptime(utc_date_str, '%Y-%m-%dT%H:%M:%SZ').replace(tzinfo=pytz.utc)
local_timezone = pytz.timezone(local_timezone_str)
local_datetime = utc_datetime.astimezone(local_timezone)
return local_datetime.strftime('%Y-%m-%d %H:%M:%S')
session = get_api_session()
flight_data = fetch_flight_data(flight, session)
dep_key = 'estimated_out' if 'estimated_out' in flight_data and flight_data['estimated_out'] else \
'actual_out' if 'actual_out' in flight_data and flight_data['actual_out'] else \
'scheduled_out'
arr_key = 'estimated_in' if 'estimated_in' in flight_data and flight_data['estimated_in'] else \
'actual_in' if 'actual_in' in flight_data and flight_data['actual_in'] else \
'scheduled_in'
flight_details = {
'flight':flight,
'source': flight_data['origin']['city'],
'destination': flight_data['destination']['city'],
'depart_time': utc_to_local(flight_data[dep_key], flight_data['origin']['timezone']),
'arrival_time': utc_to_local(flight_data[arr_key], flight_data['destination']['timezone']),
'status': flight_data['status']
}
return json.dumps(flight_details)
user_functions: Set[Callable[..., Any]] = {
get_flight_status,
}
We already imported this in our agent script with the line from user_functions import user_functions. This becomes the action tool for our agent.
Creating a Toolkit for the Agent
We will register both the tools with the toolkit that becomes available to our agent.
# Set up both tools file_search_tool = FileSearchTool(vector_store_ids=[vector_store.id]) functions = FunctionTool(user_functions) # Combine tools in a toolset toolset = ToolSet() toolset.add(functions) toolset.add(file_search_tool)
Creating an Azure AI Agent
Let’s define the agent with the instructions and a thread responsible for the final execution.
agent = project_client.agents.create_agent(
model="gpt-4o-mini",
name="unified-flight-agent",
instructions="""You are a helpful agent that can:
1. Search for information in the baggage policy document
2. Provide flight status information
Choose the appropriate tool based on the user's question.""",
toolset=toolset
)
print(f"Created agent, agent ID: {agent.id}")
# Create a thread
thread = project_client.agents.create_thread()
print(f"Created thread, thread ID: {thread.id}")
Invoke the Agent
With the agent and thread objects in place, we must create the message with the actual task and run the agent. The helper function, process_user_query, does this task.
def process_user_query(query):
message = project_client.agents.create_message(
thread_id=thread.id,
role="user",
content=query
)
print(f"Created message, message ID: {message.id}")
run = project_client.agents.create_and_process_run(thread_id=thread.id, assistant_id=agent.id)
print(f"Run status: {run.status}")
if run.status == "failed":
print(f"Run failed: {run.last_error}")
return
messages = project_client.agents.list_messages(thread_id=thread.id)
for message in messages['data']:
content = message['content']
for item in content:
if item['type'] == 'text':
print(f"Message content: {item['text']['value']}")
# Example usage
try:
# Test flight status query
process_user_query("What is the status of flight EK524?")
# Test baggage policy query
process_user_query("What is the size of checked baggage?")
finally:
# Cleanup
project_client.agents.delete_vector_store(vector_store.id)
print("Deleted vector store")
project_client.agents.delete_agent(agent.id)
print("Deleted agent")
Notice that we are sending queries that will be routed to the vector store or a function call. The LLM determines this at runtime.
When we execute, we see the below output.
I hope you found this tutorial useful in building your first AI Agent on Azure.
