2-notebooks/2-agent_service/5-agents-aisearch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4b6569b8",
   "metadata": {},
   "source": [
    "# 🏋️ AI Search + Agent Service: Fitness-Fun Example 🤸\n",
    "\n",
    "Welcome to our **AI Search + AI Agent** tutorial, where we'll:\n",
    "\n",
    "1. **Create** an Azure AI Search index with some fitness-oriented sample data\n",
    "2. **Demonstrate** how to connect that index to an Agent via the `AzureAISearchTool`\n",
    "3. **Show** how to query the Agent for health and fitness info in a fun scenario (with disclaimers!)\n",
    "\n",
    "## 🏥 Health & Fitness Disclaimer\n",
    "> **This notebook is for general demonstration and entertainment purposes, NOT a substitute for professional medical advice.**\n",
    "> Always seek the advice of certified health professionals.\n",
    "\n",
    "## Prerequisites\n",
    "1. Complete Agent basics notebook - [1-basics.ipynb](1-basics.ipynb)\n",
    "2. An **Azure AI Search** resource (formerly \"Cognitive Search\"), provisioned in your Azure AI Foundry project.\n",
    "\n",
    "## High-Level Flow\n",
    "We'll do the following:\n",
    "1. **Create** an AI Search index programmatically with sample fitness data.\n",
    "2. **Upload** documents (fitness items) to the index.\n",
    "3. **Create** an Agent that references our new index using `AzureAISearchTool`.\n",
    "4. **Run queries** to see how it fetches from the index.\n",
    " \n",
    " <img src=\"./seq-diagrams/5-ai-search.png\" width=\"30%\"/>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11b27cec",
   "metadata": {},
   "source": [
    "## 1. Create & Populate Azure AI Search Index\n",
    "We'll create a minimal index called `myfitnessindex`, containing a few example items.\n",
    "Make sure to set your environment variables for `SEARCH_ENDPOINT` and `SEARCH_API_KEY`. We'll use the `azure.search.documents.indexes` classes to manage the index schema. We'll also upload some sample data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b555e42d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required Azure libraries\n",
    "import os\n",
    "from azure.core.credentials import AzureKeyCredential  # For authentication\n",
    "from azure.search.documents.indexes import SearchIndexClient  # For managing search indexes\n",
    "from azure.search.documents.indexes.models import SearchIndex, SimpleField, SearchFieldDataType, SearchableField  # Index schema components\n",
    "from azure.search.documents import SearchClient  # For document operations (upload/search)\n",
    "from azure.identity import DefaultAzureCredential  # For Azure authentication\n",
    "from azure.ai.projects import AIProjectClient  # To access project resources\n",
    "from azure.ai.projects.models import ConnectionType  # Enum for connection types\n",
    "\n",
    "# First, initialize the AI Project client which gives us access to project resources\n",
    "# This uses DefaultAzureCredential for authentication and the project connection string\n",
    "project_client = AIProjectClient.from_connection_string(\n",
    "    credential=DefaultAzureCredential(),\n",
    "    conn_str=os.environ[\"PROJECT_CONNECTION_STRING\"]\n",
    ")\n",
    "\n",
    "# Get the Azure AI Search connection details from our project\n",
    "# This includes endpoint URL and API key needed to access the search service\n",
    "search_conn = project_client.connections.get_default(\n",
    "    connection_type=ConnectionType.AZURE_AI_SEARCH, \n",
    "    include_credentials=True\n",
    ")\n",
    "if not search_conn:\n",
    "    raise RuntimeError(\"❌ No default Azure AI Search connection found in your project.\")\n",
    "\n",
    "# Name of our search index - this is where our fitness data will be stored\n",
    "index_name = \"myfitnessindex\"\n",
    "\n",
    "try:\n",
    "    # Create a SearchIndexClient - this is used for managing the index itself (create/update/delete)\n",
    "    credential = AzureKeyCredential(search_conn.key)\n",
    "    index_client = SearchIndexClient(endpoint=search_conn.endpoint_url, credential=credential)\n",
    "    print(\"✅ Created SearchIndexClient from project_client connection\")\n",
    "    \n",
    "    # Create a SearchClient - this is used for document operations (upload/search/delete documents)\n",
    "    # We'll use this later to add our fitness items to the index\n",
    "    search_client = SearchClient(\n",
    "        endpoint=search_conn.endpoint_url,\n",
    "        index_name=index_name,\n",
    "        credential=credential\n",
    "    )\n",
    "    print(\"✅ Created SearchClient for document operations\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(f\"❌ Error creating search clients: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8c418b5",
   "metadata": {},
   "source": [
    "**Define the index** schema with a `FitnessItemID` key and a few fields to store product info.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eed5e598",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_fitness_index():\n",
    "    # Define the fields (columns) for our search index\n",
    "    # Each field has specific attributes that control how it can be used in searches:\n",
    "    fields = [\n",
    "        # Primary key field - must be unique for each document\n",
    "        SimpleField(name=\"FitnessItemID\", type=SearchFieldDataType.String, key=True),\n",
    "        \n",
    "        # Name field - SearchableField means we can do full-text search on it\n",
    "        # filterable=True lets us filter results by name\n",
    "        SearchableField(name=\"Name\", type=SearchFieldDataType.String, filterable=True),\n",
    "        \n",
    "        # Category field - SearchableField for text search\n",
    "        # filterable=True lets us filter by category\n",
    "        # facetable=True enables category grouping in results\n",
    "        SearchableField(name=\"Category\", type=SearchFieldDataType.String, filterable=True, facetable=True),\n",
    "        \n",
    "        # Price field - SimpleField for numeric values\n",
    "        # filterable=True enables price range filters\n",
    "        # sortable=True lets us sort by price\n",
    "        # facetable=True enables price range grouping\n",
    "        SimpleField(name=\"Price\", type=SearchFieldDataType.Double, filterable=True, sortable=True, facetable=True),\n",
    "        \n",
    "        # Description field - SearchableField for full-text search on product descriptions\n",
    "        SearchableField(name=\"Description\", type=SearchFieldDataType.String)\n",
    "    ]\n",
    "\n",
    "    # Create an index definition with our fields\n",
    "    index = SearchIndex(name=index_name, fields=fields)\n",
    "\n",
    "    # Check if index already exists - if so, delete it to start fresh\n",
    "    # This is useful during development but be careful in production!\n",
    "    if index_name in [x.name for x in index_client.list_indexes()]:\n",
    "        index_client.delete_index(index_name)\n",
    "        print(f\"🗑️ Deleted existing index: {index_name}\")\n",
    "\n",
    "    # Create the new index with our schema\n",
    "    created = index_client.create_index(index)\n",
    "    print(f\"🎉 Created index: {created.name}\")\n",
    "\n",
    "# Execute the function to create our search index\n",
    "create_fitness_index()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17035c71",
   "metadata": {},
   "source": [
    "**Upload some sample documents** to `myfitnessindex`. We'll add a few items for demonstration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9af38671",
   "metadata": {},
   "outputs": [],
   "source": [
    "def upload_fitness_docs():\n",
    "    # Create a SearchClient to interact with our search index\n",
    "    # This uses the connection details (endpoint, key) we configured earlier\n",
    "    search_client = SearchClient(\n",
    "        endpoint=search_conn.endpoint_url,\n",
    "        index_name=index_name,\n",
    "        credential=AzureKeyCredential(search_conn.key)\n",
    "    )\n",
    "\n",
    "    # Define sample documents that match our index schema\n",
    "    # Each document must have:\n",
    "    # - FitnessItemID (unique identifier)\n",
    "    # - Name (searchable product name) \n",
    "    # - Category (searchable and facetable for filtering/grouping)\n",
    "    # - Price (numeric field for sorting and filtering)\n",
    "    # - Description (searchable product details)\n",
    "    sample_docs = [\n",
    "        {\n",
    "            \"FitnessItemID\": \"1\",\n",
    "            \"Name\": \"Adjustable Dumbbell\",\n",
    "            \"Category\": \"Strength\", \n",
    "            \"Price\": 59.99,\n",
    "            \"Description\": \"A compact, adjustable weight for targeted muscle workouts.\"\n",
    "        },\n",
    "        {\n",
    "            \"FitnessItemID\": \"2\",\n",
    "            \"Name\": \"Yoga Mat\",\n",
    "            \"Category\": \"Flexibility\",\n",
    "            \"Price\": 25.0,\n",
    "            \"Description\": \"Non-slip mat designed for yoga, Pilates, and other exercises.\"\n",
    "        },\n",
    "        {\n",
    "            \"FitnessItemID\": \"3\",\n",
    "            \"Name\": \"Treadmill\",\n",
    "            \"Category\": \"Cardio\",\n",
    "            \"Price\": 499.0,\n",
    "            \"Description\": \"A sturdy treadmill with adjustable speed and incline settings.\"\n",
    "        },\n",
    "        {\n",
    "            \"FitnessItemID\": \"4\",\n",
    "            \"Name\": \"Resistance Bands\",\n",
    "            \"Category\": \"Strength\",\n",
    "            \"Price\": 15.0,\n",
    "            \"Description\": \"Set of colorful bands for light to moderate resistance workouts.\"\n",
    "        }\n",
    "    ]\n",
    "\n",
    "    # Upload all documents to the search index in a single batch operation\n",
    "    # The search service will index these documents, making them searchable\n",
    "    # based on the field types we defined in our index schema\n",
    "    result = search_client.upload_documents(documents=sample_docs)\n",
    "    print(f\"🚀 Upload result: {result}\")\n",
    "\n",
    "# Call the function to upload the documents\n",
    "upload_fitness_docs()\n",
    "print(\"✅ Documents uploaded to search index\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d95b1386",
   "metadata": {},
   "source": [
    "### Verify the documents via a basic query\n",
    "Let's do a quick search for **Strength** items.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2af93084",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's verify our index by performing a basic search\n",
    "# 1. First create a SearchClient using our connection details\n",
    "#    - endpoint_url: The URL of our search service\n",
    "#    - index_name: The name of the index we created earlier\n",
    "#    - key: The admin key to authenticate our requests\n",
    "search_client = SearchClient(\n",
    "    endpoint=search_conn.endpoint_url,\n",
    "    index_name=index_name,\n",
    "    credential=AzureKeyCredential(search_conn.key)\n",
    ")\n",
    "\n",
    "# 2. Perform a simple search query:\n",
    "#    - search_text=\"Strength\": Look for documents containing \"Strength\"\n",
    "#    - filter=None: No additional filtering\n",
    "#    - top=10: Return up to 10 matching documents\n",
    "results = search_client.search(search_text=\"Strength\", filter=None, top=10)\n",
    "\n",
    "# 3. Print each matching document\n",
    "print(\"🔍 Search results for 'Strength':\")\n",
    "print(\"-\" * 50)\n",
    "found_items = False\n",
    "for doc in results:\n",
    "    found_items = True\n",
    "    # The doc is already a dictionary, no need for to_dict()\n",
    "    print(f\"Name: {doc['Name']}\")\n",
    "    print(f\"Category: {doc['Category']}\")\n",
    "    print(f\"Price: ${doc['Price']:.2f}\")\n",
    "    print(f\"Description: {doc['Description']}\")\n",
    "    print(\"-\" * 50)\n",
    "\n",
    "if not found_items:\n",
    "    print(\"No matching items found.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed67b2f4",
   "metadata": {},
   "source": [
    "## 2. Create Agent With AI Search Tool\n",
    "We'll create a new agent and attach an `AzureAISearchTool` referencing **myfitnessindex**.\n",
    "In your environment, you need:\n",
    "- `PROJECT_CONNECTION_STRING` - from your AI Foundry project overview\n",
    "- `MODEL_DEPLOYMENT_NAME` - from the deployed model name\n",
    "\n",
    "Let's initialize the `AIProjectClient` with `DefaultAzureCredential`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c31d0e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries:\n",
    "# - os: For accessing environment variables\n",
    "# - DefaultAzureCredential: Azure's authentication mechanism\n",
    "# - AIProjectClient: Main client for interacting with AI Projects\n",
    "# - AzureAISearchTool & ConnectionType: Used to configure search capabilities\n",
    "import os\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.ai.projects import AIProjectClient\n",
    "from azure.ai.projects.models import AzureAISearchTool, ConnectionType\n",
    "\n",
    "# Initialize the AI Project Client which we'll use to:\n",
    "# 1. Connect to our Azure AI project\n",
    "# 2. Create agents with search capabilities\n",
    "# 3. Manage project resources\n",
    "try:\n",
    "    project_client = AIProjectClient.from_connection_string(\n",
    "        # Use Azure's default authentication method\n",
    "        credential=DefaultAzureCredential(),\n",
    "        # Connect using the project connection string from environment variables\n",
    "        conn_str=os.environ[\"PROJECT_CONNECTION_STRING\"],\n",
    "    )\n",
    "    print(\"✅ Successfully initialized AIProjectClient\")\n",
    "except Exception as e:\n",
    "    print(f\"❌ Error initializing project client: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4988b6c",
   "metadata": {},
   "source": [
    "### Find (or create) the Azure AI Search connection in your Foundry project\n",
    "We'll now use `project_client.connections.get_default(...)` to retrieve the default Azure AI Search connection, **including** credentials.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9de13ece",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try to get the default Azure AI Search connection from our project\n",
    "# - Azure AI Search (formerly Cognitive Search) is a cloud search service that helps \n",
    "#   us add search capabilities to our applications\n",
    "# - The connection contains endpoint and credential information needed to access the search service\n",
    "search_conn = project_client.connections.get_default(\n",
    "    # Specify we want an Azure AI Search connection type\n",
    "    connection_type=ConnectionType.AZURE_AI_SEARCH,\n",
    "    # include_credentials=True means we'll get the full connection info including auth keys\n",
    "    include_credentials=True\n",
    ")\n",
    "\n",
    "# Check if we found a connection\n",
    "if not search_conn:\n",
    "    print(\"❌ No default Azure AI Search connection found in your project.\")\n",
    "else:\n",
    "    # If found, print the connection details\n",
    "    # - The connection ID is a unique identifier for this connection in our project\n",
    "    # - The endpoint URL is where our search service is hosted\n",
    "    print(f\"Found default Azure AI Search connection ID: {search_conn.id}\")\n",
    "    print(f\"Index endpoint: {search_conn.endpoint_url}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e921f6c7",
   "metadata": {},
   "source": [
    "### Create the Agent with `AzureAISearchTool`\n",
    "We'll attach the tool, specifying the index name we created.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e2fe256",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the model deployment name from environment variables\n",
    "# This is the Azure OpenAI model we'll use for our agent\n",
    "model_name = os.environ.get(\"MODEL_DEPLOYMENT_NAME\")\n",
    "agent = None\n",
    "\n",
    "if search_conn:\n",
    "    # Create an Azure AI Search tool that will allow our agent to search the fitness equipment index\n",
    "    # - The tool needs the connection ID we got earlier to authenticate\n",
    "    # - index_name specifies which search index to query (we created myfitnessindex earlier)\n",
    "    ai_search_tool = AzureAISearchTool(\n",
    "        index_connection_id=search_conn.id,\n",
    "        index_name=index_name\n",
    "    )\n",
    "\n",
    "    # Create an AI agent that can understand natural language and search our index\n",
    "    # - The agent uses our Azure OpenAI model for natural language understanding\n",
    "    # - We give it instructions to act as a fitness shopping assistant\n",
    "    # - We attach the search tool so it can look up products\n",
    "    # - tool_resources provides the connection details the tool needs\n",
    "    agent = project_client.agents.create_agent(\n",
    "        model=model_name,\n",
    "        name=\"fitness-agent-search\",\n",
    "        instructions=\"\"\"\n",
    "        You are a Fitness Shopping Assistant. You help users find items, but always disclaim not to provide medical advice.\n",
    "        \"\"\",\n",
    "        tools=ai_search_tool.definitions,\n",
    "        tool_resources=ai_search_tool.resources,\n",
    "        headers={\"x-ms-enable-preview\": \"true\"},  # Enable preview features\n",
    "    )\n",
    "    print(f\"🎉 Created agent, ID: {agent.id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "704445fc",
   "metadata": {},
   "source": [
    "## 3. Run a Conversation with the Agent\n",
    "We'll open a new thread, post a question, and let the agent search the index for relevant items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ba65c9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_agent_query(question: str):\n",
    "    # Step 1: Create a new conversation thread\n",
    "    # In Azure AI Agent service, conversations happen in threads, similar to chat conversations\n",
    "    # Each thread can contain multiple back-and-forth messages\n",
    "    thread = project_client.agents.create_thread()\n",
    "    print(f\"📝 Created thread, ID: {thread.id}\")\n",
    "\n",
    "    # Step 2: Add the user's question as a message in the thread\n",
    "    # Messages have roles (\"user\" or \"assistant\") and content (the actual text)\n",
    "    message = project_client.agents.create_message(\n",
    "        thread_id=thread.id,\n",
    "        role=\"user\",\n",
    "        content=question\n",
    "    )\n",
    "    print(f\"💬 Created user message, ID: {message.id}\")\n",
    "\n",
    "    # Step 3: Create and start an agent run\n",
    "    # This tells the agent to:\n",
    "    # - Read the user's message\n",
    "    # - Use its AI Search tool to find relevant products\n",
    "    # - Generate a helpful response\n",
    "    run = project_client.agents.create_and_process_run(\n",
    "        thread_id=thread.id,\n",
    "        agent_id=agent.id\n",
    "    )\n",
    "    print(f\"🤖 Agent run status: {run.status}\")\n",
    "\n",
    "    # Check for any errors during the agent's processing\n",
    "    if run.last_error:\n",
    "        print(\"⚠️ Run error:\", run.last_error)\n",
    "\n",
    "    # Step 4: Get the agent's response\n",
    "    # Retrieve all messages and find the most recent assistant response\n",
    "    # The response might contain multiple content blocks (text, images, etc.)\n",
    "    msg_list = project_client.agents.list_messages(thread_id=thread.id)\n",
    "    for m in reversed(msg_list.data):\n",
    "        if m.role == \"assistant\" and m.content:\n",
    "            print(\"\\nAssistant says:\")\n",
    "            for c in m.content:\n",
    "                if hasattr(c, \"text\"):\n",
    "                    print(c.text.value)\n",
    "            break\n",
    "\n",
    "# Try out our agent with two example queries:\n",
    "# 1. A general question about strength training equipment\n",
    "# 2. A specific request for cardio equipment with a price constraint\n",
    "if agent:\n",
    "    run_agent_query(\"Which items are good for strength training?\")\n",
    "    run_agent_query(\"I need something for cardio under $300, any suggestions?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0292c370",
   "metadata": {},
   "source": [
    "## 4. Cleanup\n",
    "We'll clean up the agent. (In production, you might want to keep it!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "285fdc95",
   "metadata": {},
   "outputs": [],
   "source": [
    "if agent:\n",
    "    project_client.agents.delete_agent(agent.id)\n",
    "    print(\"🗑️ Deleted agent\")\n",
    "\n",
    "index_client.delete_index(index_name)\n",
    "print(f\"🗑️ Deleted index {index_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcf584e3",
   "metadata": {},
   "source": [
    "# 🎉 Congrats!\n",
    "You've successfully:\n",
    "1. **Created** an Azure AI Search index programmatically.\n",
    "2. **Populated** it with sample fitness data.\n",
    "3. **Created** an Agent that queries the index using `AzureAISearchTool`.\n",
    "4. **Asked** the agent for item recommendations.\n",
    "\n",
    "Continue exploring how to integrate **OpenTelemetry** or the `azure-ai-evaluation` library for advanced tracing and evaluation capabilities. Have fun, and stay fit! 🏆"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}