notebook/agentchat_azr_ai_search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Assistants with Azure Cognitive Search and Azure Identity\n",
    "\n",
    "This notebook demonstrates the use of Assistant Agents in conjunction with Azure Cognitive Search and Azure Identity. Assistant Agents use tools that interact with Azure Cognitive Search to extract pertinent data.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "Before running this notebook, please ensure the following prerequisites are met:\n",
    " \n",
    "\n",
    "### Dependencies\n",
    "1. **Autogen**\n",
    "2. **Azure SDK**\n",
    "3. **Cognitive Search**/**AI Search**\n",
    "\n",
    "If you have AI search enabled in your Azure Portal, you can use the following code to create an assistant agent that can search Azure Cognitive Search.\n",
    "\n",
    "**AI search setup details:**\n",
    "- Documentation:   \n",
    "    - Create search service: https://learn.microsoft.com/en-us/azure/search/search-create-service-portal \n",
    "    - Search index:  https://learn.microsoft.com/en-us/azure/search/search-how-to-create-search-index?tabs=portal \n",
    "    hybrid search: https://learn.microsoft.com/en-us/azure/search/hybrid-search-how-to-query\n",
    "\n",
    "- Youtube walkthrough: https://www.youtube.com/watch?v=6Zfuw-UJZ7k\n",
    "\n",
    "\n",
    "### Install Azure CLI\n",
    "This notebook requires the Azure CLI for authentication purposes. Follow these steps to install and configure it:\n",
    "\n",
    "1. **Download and Install Azure CLI**:\n",
    "   - Visit the [Azure CLI installation page](https://docs.microsoft.com/en-us/cli/azure/install-azure-cli) and follow the instructions for your operating system.\n",
    "   - Mac users can install Azure CLI using Homebrew with the command `brew install azure-cli`   \n",
    "\n",
    "2. **Verify Installation**:\n",
    "   - In the below cell execute `az --version` to check if Azure CLI is installed correctly.\n",
    "\n",
    "4. **Login to Azure**:\n",
    "   - In the below cell execute `az login` to log into your Azure account. This step is necessary as the notebook uses `AzureCliCredential` which retrieves the token based on the Azure account currently logged in.\n",
    "\n",
    "### Check Azure CLI Installation\n",
    "Run the cell below to check if Azure CLI is installed and properly configured on your system."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check Azure CLI Installation and Login Status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check Azure CLI installation and login status\n",
    "# !az --version\n",
    "# !az login"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install required packages\n",
    "Run the cell below to install the required packages for this notebook.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install pyautogen==0.2.16\n",
    "!pip3 install python-dotenv==1.0.1\n",
    "!pip3 install pyautogen[graph]>=0.2.11\n",
    "!pip3 install azure-search-documents==11.4.0b8\n",
    "!pip3 install azure-identity==1.12.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next you will import the required packages for this notebook.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.search.documents import SearchClient\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "import autogen\n",
    "from autogen import AssistantAgent, UserProxyAgent, register_function\n",
    "from autogen.cache import Cache\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "# Import Cognitive Search index ENV\n",
    "AZURE_SEARCH_SERVICE = os.getenv(\"AZURE_SEARCH_SERVICE\")\n",
    "AZURE_SEARCH_INDEX = os.getenv(\"AZURE_SEARCH_INDEX\")\n",
    "AZURE_SEARCH_KEY = os.getenv(\"AZURE_SEARCH_KEY\")\n",
    "AZURE_SEARCH_API_VERSION = os.getenv(\"AZURE_SEARCH_API_VERSION\")\n",
    "AZURE_SEARCH_SEMANTIC_SEARCH_CONFIG = os.getenv(\"AZURE_SEARCH_SEMANTIC_SEARCH_CONFIG\")\n",
    "AZURE_SEARCH_SERVICE_ENDPOINT = os.getenv(\"AZURE_SEARCH_SERVICE_ENDPOINT\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, you need to authenticate and create a `SearchClient` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "credential = DefaultAzureCredential()\n",
    "endpoint = AZURE_SEARCH_SERVICE_ENDPOINT\n",
    "\n",
    "from azure.identity import AzureCliCredential\n",
    "\n",
    "credential = AzureCliCredential()\n",
    "token = credential.get_token(\"https://cognitiveservices.azure.com/.default\")\n",
    "\n",
    "print(\"TOKEN\", token.token)\n",
    "\n",
    "client = SearchClient(endpoint=endpoint, index_name=\"test-index\", credential=credential)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Then, load the configuration list and define the configuration for the `AssistantAgent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_list = autogen.config_list_from_json(\n",
    "    env_or_file=\"OAI_CONFIG_LIST\",\n",
    ")\n",
    "\n",
    "gpt4_config = {\n",
    "    \"cache_seed\": 42,\n",
    "    \"temperature\": 0,\n",
    "    \"config_list\": config_list,\n",
    "    \"timeout\": 120,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Define your tool function `search` that will interact with the Azure Cognitive Search service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search(query: str):\n",
    "    payload = json.dumps({\n",
    "        \"search\": query,\n",
    "        \"vectorQueries\": [{\"kind\": \"text\", \"text\": query, \"k\": 5, \"fields\": \"vector\"}],\n",
    "        \"queryType\": \"semantic\",\n",
    "        \"semanticConfiguration\": AZURE_SEARCH_SEMANTIC_SEARCH_CONFIG,\n",
    "        \"captions\": \"extractive\",\n",
    "        \"answers\": \"extractive|count-3\",\n",
    "        \"queryLanguage\": \"en-US\",\n",
    "    })\n",
    "\n",
    "    response = list(client.search(payload))\n",
    "\n",
    "    output = []\n",
    "    for result in response:\n",
    "        result.pop(\"titleVector\")\n",
    "        result.pop(\"contentVector\")\n",
    "        output.append(result)\n",
    "\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Define the `AssistantAgent` and `UserProxyAgent` instances, and register the `search` function to them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cog_search = AssistantAgent(\n",
    "    name=\"COGSearch\",\n",
    "    system_message=\"You are a helpful AI assistant. \"\n",
    "    \"You can help with Azure Cognitive Search.\"\n",
    "    \"Return 'TERMINATE' when the task is done.\",\n",
    "    llm_config=gpt4_config,\n",
    ")\n",
    "\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"User\",\n",
    "    llm_config=False,\n",
    "    is_termination_msg=lambda msg: msg.get(\"content\") is not None and \"TERMINATE\" in msg[\"content\"],\n",
    "    human_input_mode=\"NEVER\",\n",
    ")\n",
    "\n",
    "register_function(\n",
    "    search,\n",
    "    caller=cog_search,\n",
    "    executor=user_proxy,\n",
    "    name=\"search\",\n",
    "    description=\"A tool for searching the Cognitive Search index\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, initiate a chat."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "\n",
    "    async def main():\n",
    "        with Cache.disk() as cache:\n",
    "            await user_proxy.a_initiate_chat(\n",
    "                cog_search,\n",
    "                message=\"Search for 'What is Azure?' in the 'test-index' index\",\n",
    "                cache=cache,\n",
    "            )\n",
    "\n",
    "    await main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "This notebook demonstrates the use of Assistant Agents in conjunction with Azure Cognitive Search and Azure Identity",
   "tags": [
    "integration",
    "RAG",
    "Azure Identity",
    "Azure AI Search"
   ]
  },
  "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.12.3"
  },
  "skip_test": "This requires Azure AI Search to be enabled and creds for AI Search from Azure Portal"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}