enterprise-streaming-agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Azure AI Agent Service Enterprise Demo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Necessary Libraries\n",
    "In this cell, we import all the libraries and modules required for the project.\n",
    "This includes Azure AI SDKs, Gradio for UI, and custom functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "from datetime import datetime as pydatetime\n",
    "from typing import Any, List, Dict\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# (Optional) Gradio app for UI\n",
    "import gradio as gr\n",
    "from gradio import ChatMessage\n",
    "\n",
    "# Azure AI Projects\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.ai.projects import AIProjectClient\n",
    "from azure.ai.projects.models import (\n",
    "    AgentEventHandler,\n",
    "    RunStep,\n",
    "    RunStepDeltaChunk,\n",
    "    ThreadMessage,\n",
    "    ThreadRun,\n",
    "    MessageDeltaChunk,\n",
    "    BingGroundingTool,\n",
    "    FilePurpose,\n",
    "    FileSearchTool,\n",
    "    FunctionTool,\n",
    "    ToolSet\n",
    ")\n",
    "\n",
    "# Your custom Python functions (for \"fetch_weather\",\"fetch_stock_price\",\"send_email\",\"fetch_datetime\", etc.)\n",
    "from enterprise_functions import enterprise_fns\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Client and Load Azure AI Foundry\n",
    "Here, we initialize the Azure AI client using DefaultAzureCredential.\n",
    "This allows us to authenticate and connect to the Azure AI service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "credential = DefaultAzureCredential()\n",
    "project_client = AIProjectClient.from_connection_string(\n",
    "    credential=credential,\n",
    "    conn_str=os.environ[\"PROJECT_CONNECTION_STRING\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set Up Tools (BingGroundingTool, FileSearchTool)\n",
    "In this step, we configure tools such as `BingGroundingTool` and `FileSearchTool`.\n",
    "We check for existing connections and create or reuse vector stores for document search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    bing_connection = project_client.connections.get(connection_name=os.environ[\"BING_CONNECTION_NAME\"])\n",
    "    conn_id = bing_connection.id\n",
    "    bing_tool = BingGroundingTool(connection_id=conn_id)\n",
    "    print(\"bing > connected\")\n",
    "except Exception:\n",
    "    bing_tool = None\n",
    "    print(\"bing failed > no connection found or permission issue\")\n",
    "\n",
    "FOLDER_NAME = \"enterprise-data\"\n",
    "VECTOR_STORE_NAME = \"hr-policy-vector-store\"\n",
    "all_vector_stores = project_client.agents.list_vector_stores().data\n",
    "existing_vector_store = next(\n",
    "    (store for store in all_vector_stores if store.name == VECTOR_STORE_NAME),\n",
    "    None\n",
    ")\n",
    "\n",
    "vector_store_id = None\n",
    "if existing_vector_store:\n",
    "    vector_store_id = existing_vector_store.id\n",
    "    print(f\"reusing vector store > {existing_vector_store.name} (id: {existing_vector_store.id})\")\n",
    "else:\n",
    "    # If you have local docs to upload\n",
    "    import os\n",
    "    if os.path.isdir(FOLDER_NAME):\n",
    "        file_ids = []\n",
    "        for file_name in os.listdir(FOLDER_NAME):\n",
    "            file_path = os.path.join(FOLDER_NAME, file_name)\n",
    "            if os.path.isfile(file_path):\n",
    "                print(f\"uploading > {file_name}\")\n",
    "                uploaded_file = project_client.agents.upload_file_and_poll(\n",
    "                    file_path=file_path,\n",
    "                    purpose=FilePurpose.AGENTS\n",
    "                )\n",
    "                file_ids.append(uploaded_file.id)\n",
    "\n",
    "        if file_ids:\n",
    "            print(f\"creating vector store > from {len(file_ids)} files.\")\n",
    "            vector_store = project_client.agents.create_vector_store_and_poll(\n",
    "                file_ids=file_ids,\n",
    "                name=VECTOR_STORE_NAME\n",
    "            )\n",
    "            vector_store_id = vector_store.id\n",
    "            print(f\"created > {vector_store.name} (id: {vector_store_id})\")\n",
    "\n",
    "file_search_tool = None\n",
    "if vector_store_id:\n",
    "    file_search_tool = FileSearchTool(vector_store_ids=[vector_store_id])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine All Tools into a ToolSet\n",
    "This step creates a custom `ToolSet` that includes all the tools configured earlier.\n",
    "It also adds a `LoggingToolSet` subclass to log the inputs and outputs of function calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LoggingToolSet(ToolSet):\n",
    "    def execute_tool_calls(self, tool_calls: List[Any]) -> List[dict]:\n",
    "        \"\"\"\n",
    "        Execute the upstream calls, printing only two lines per function:\n",
    "        1) The function name + its input arguments\n",
    "        2) The function name + its output result\n",
    "        \"\"\"\n",
    "\n",
    "        # For each function call, print the input arguments\n",
    "        for c in tool_calls:\n",
    "            if hasattr(c, \"function\") and c.function:\n",
    "                fn_name = c.function.name\n",
    "                fn_args = c.function.arguments\n",
    "                print(f\"{fn_name} inputs > {fn_args} (id:{c.id})\")\n",
    "\n",
    "        # Execute the tool calls (superclass logic)\n",
    "        raw_outputs = super().execute_tool_calls(tool_calls)\n",
    "\n",
    "        # Print the output of each function call\n",
    "        for item in raw_outputs:\n",
    "            print(f\"output > {item['output']}\")\n",
    "\n",
    "        return raw_outputs\n",
    "\n",
    "custom_functions = FunctionTool(enterprise_fns)\n",
    "\n",
    "toolset = LoggingToolSet()\n",
    "if bing_tool:\n",
    "    toolset.add(bing_tool)\n",
    "if file_search_tool:\n",
    "    toolset.add(file_search_tool)\n",
    "toolset.add(custom_functions)\n",
    "\n",
    "for tool in toolset._tools:\n",
    "    tool_name = tool.__class__.__name__\n",
    "    print(f\"tool > {tool_name}\")\n",
    "    for definition in tool.definitions:\n",
    "        if hasattr(definition, \"function\"):\n",
    "            fn = definition.function\n",
    "            print(f\"{fn.name} > {fn.description}\")\n",
    "        else:\n",
    "            pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Optional) Direct Azure AI Search Integration\n",
    "Skip this cell if you're using the default File Search Tool vector store approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (Optional) Direct Azure AI Search Integration\n",
    "# If the default File Search Tool is available, we add it and skip the Azure AI Search integration.\n",
    "if any(tool.__class__.__name__ == \"FileSearchTool\" for tool in toolset._tools):\n",
    "    print(\"file_search tool exists > skipping ai_search tool add\")\n",
    "else:\n",
    "    try:\n",
    "        # Get the connection ID for your Azure AI Search resource\n",
    "        connections = project_client.connections.list()\n",
    "        conn_id = next(\n",
    "            c.id for c in connections if c.name == os.environ.get(\"AZURE_SEARCH_CONNECTION_NAME\")\n",
    "        )\n",
    "        \n",
    "        # Initialize Azure AI Search tool for direct index access\n",
    "        from azure.ai.projects.models import AzureAISearchTool\n",
    "        search_tool = AzureAISearchTool(\n",
    "            index_connection_id=conn_id,\n",
    "            index_name=os.environ.get(\"AZURE_SEARCH_INDEX_NAME\")\n",
    "        )\n",
    "        \n",
    "        # Add the Azure AI Search tool to our toolset\n",
    "        toolset.add(search_tool)\n",
    "        print(\"azure ai search > connected directly to index\")\n",
    "        \n",
    "        # Verify the tool was added by iterating through the toolset\n",
    "        for tool in toolset._tools:\n",
    "            tool_name = tool.__class__.__name__\n",
    "            print(f\"tool > {tool_name}\")\n",
    "            for definition in tool.definitions:\n",
    "                if hasattr(definition, \"function\"):\n",
    "                    fn = definition.function\n",
    "                    print(f\"{fn.name} > {fn.description}\")\n",
    "                else:\n",
    "                    pass\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"azure ai search > skipped (no connection configured): {str(e)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create or Reuse the Enterprise Agent\n",
    "In this step, we create a new enterprise agent or reuse an existing one.\n",
    "The agent is configured with a model, instructions, and the toolset from the previous step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AGENT_NAME = \"my-enterprise-agent\"\n",
    "found_agent = None\n",
    "all_agents_list = project_client.agents.list_agents().data\n",
    "for a in all_agents_list:\n",
    "    if a.name == AGENT_NAME:\n",
    "        found_agent = a\n",
    "        break\n",
    "\n",
    "model_name = os.environ.get(\"MODEL_DEPLOYMENT_NAME\", \"gpt-4o\")\n",
    "instructions = (\n",
    "    \"You are a helpful enterprise assistant at Contoso. \"\n",
    "    f\"Today's date is {pydatetime.now().strftime('%A, %b %d, %Y, %I:%M %p')}. \"\n",
    "    \"You have access to hr documents in file_search, the grounding engine from bing and custom python functions like fetch_weather, \"\n",
    "    \"fetch_stock_price, send_email, etc. Provide well-structured, concise, and professional answers.\"\n",
    ")\n",
    "\n",
    "if found_agent:\n",
    "    # Update the existing agent to use new tools\n",
    "    agent = project_client.agents.update_agent(\n",
    "        assistant_id=found_agent.id,\n",
    "        model=found_agent.model,\n",
    "        instructions=found_agent.instructions,\n",
    "        toolset=toolset,\n",
    "    )\n",
    "    print(f\"reusing agent > {agent.name} (id: {agent.id})\")\n",
    "else:\n",
    "    agent = project_client.agents.create_agent(\n",
    "        model=model_name,\n",
    "        name=AGENT_NAME,\n",
    "        instructions=instructions,\n",
    "        toolset=toolset\n",
    "    )\n",
    "    print(f\"creating agent > {agent.name} (id: {agent.id})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a Conversation Thread\n",
    "In this step, we create a new conversation thread for the enterprise agent.\n",
    "Threads are used to manage and track conversations with the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thread = project_client.agents.create_thread()\n",
    "print(f\"thread > created (id: {thread.id})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a Custom Event Handler\n",
    "Here, we define a custom event handler to manage logs and outputs for debugging.\n",
    "This handler will capture and display real-time events during the agent's operation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyEventHandler(AgentEventHandler):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self._current_message_id = None\n",
    "        self._accumulated_text = \"\"\n",
    "\n",
    "    def on_message_delta(self, delta: MessageDeltaChunk) -> None:\n",
    "        # If a new message id, start fresh\n",
    "        if delta.id != self._current_message_id:\n",
    "            # First, if we had an old message that wasn't completed, finish that line\n",
    "            if self._current_message_id is not None:\n",
    "                print()  # move to a new line\n",
    "            \n",
    "            self._current_message_id = delta.id\n",
    "            self._accumulated_text = \"\"\n",
    "            print(\"\\nassistant > \", end=\"\")  # prefix for new message\n",
    "\n",
    "        # Accumulate partial text\n",
    "        partial_text = \"\"\n",
    "        if delta.delta.content:\n",
    "            for chunk in delta.delta.content:\n",
    "                partial_text += chunk.text.get(\"value\", \"\")\n",
    "        self._accumulated_text += partial_text\n",
    "\n",
    "        # Print partial text with no newline\n",
    "        print(partial_text, end=\"\", flush=True)\n",
    "\n",
    "    def on_thread_message(self, message: ThreadMessage) -> None:\n",
    "        # When the assistant's entire message is \"completed\", print a final newline\n",
    "        if message.status == \"completed\" and message.role == \"assistant\":\n",
    "            print()  # done with this line\n",
    "            self._current_message_id = None\n",
    "            self._accumulated_text = \"\"\n",
    "        else:\n",
    "            # For other roles or statuses, you can log if you like:\n",
    "            print(f\"{message.status.name.lower()} (id: {message.id})\")\n",
    "\n",
    "    def on_thread_run(self, run: ThreadRun) -> None:\n",
    "        print(f\"status > {run.status.name.lower()}\")\n",
    "        if run.status == \"failed\":\n",
    "            print(f\"error > {run.last_error}\")\n",
    "\n",
    "    def on_run_step(self, step: RunStep) -> None:\n",
    "        print(f\"{step.type.name.lower()} > {step.status.name.lower()}\")\n",
    "\n",
    "    def on_run_step_delta(self, delta: RunStepDeltaChunk) -> None:\n",
    "        # If partial tool calls come in, we log them\n",
    "        if delta.delta.step_details and delta.delta.step_details.tool_calls:\n",
    "            for tcall in delta.delta.step_details.tool_calls:\n",
    "                if getattr(tcall, \"function\", None):\n",
    "                    if tcall.function.name is not None:\n",
    "                        print(f\"tool call > {tcall.function.name}\")\n",
    "\n",
    "    def on_unhandled_event(self, event_type: str, event_data):\n",
    "        print(f\"unhandled > {event_type} > {event_data}\")\n",
    "\n",
    "    def on_error(self, data: str) -> None:\n",
    "        print(f\"error > {data}\")\n",
    "\n",
    "    def on_done(self) -> None:\n",
    "        print(\"done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implement the Main Chat Functions\n",
    "These functions define how user messages and tool interactions are processed.\n",
    "It uses the agent's thread to handle conversations and streams partial responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_bing_query(request_url: str) -> str:\n",
    "    \"\"\"\n",
    "    Extract the query string from something like:\n",
    "      https://api.bing.microsoft.com/v7.0/search?q=\"latest news about Microsoft January 2025\"\n",
    "    Returns: latest news about Microsoft January 2025\n",
    "    \"\"\"\n",
    "    match = re.search(r'q=\"([^\"]+)\"', request_url)\n",
    "    if match:\n",
    "        return match.group(1)\n",
    "    # If no match, fall back to entire request_url\n",
    "    return request_url\n",
    "\n",
    "def convert_dict_to_chatmessage(msg: dict) -> ChatMessage:\n",
    "    \"\"\"\n",
    "    Convert a legacy dict-based message to a gr.ChatMessage.\n",
    "    Uses the 'metadata' sub-dict if present.\n",
    "    \"\"\"\n",
    "    return ChatMessage(\n",
    "        role=msg[\"role\"],\n",
    "        content=msg[\"content\"],\n",
    "        metadata=msg.get(\"metadata\", None)\n",
    "    )\n",
    "\n",
    "def azure_enterprise_chat(user_message: str, history: List[dict]):\n",
    "    \"\"\"\n",
    "    Accumulates partial function arguments into ChatMessage['content'], sets the\n",
    "    corresponding tool bubble status from \"pending\" to \"done\" on completion,\n",
    "    and also handles non-function calls like bing_grounding or file_search by appending a\n",
    "    \"pending\" bubble. Then it moves them to \"done\" once tool calls complete.\n",
    "\n",
    "    This function returns a list of ChatMessage objects directly (no dict conversion).\n",
    "    Your Gradio Chatbot should be type=\"messages\" to handle them properly.\n",
    "    \"\"\"\n",
    "\n",
    "    # Convert existing history from dict to ChatMessage\n",
    "    conversation = []\n",
    "    for msg_dict in history:\n",
    "        conversation.append(convert_dict_to_chatmessage(msg_dict))\n",
    "\n",
    "    # Append the user's new message\n",
    "    conversation.append(ChatMessage(role=\"user\", content=user_message))\n",
    "\n",
    "    # Immediately yield two outputs to clear the textbox\n",
    "    yield conversation, \"\"\n",
    "\n",
    "    # Post user message to the thread (for your back-end logic)\n",
    "    project_client.agents.create_message(\n",
    "        thread_id=thread.id,\n",
    "        role=\"user\",\n",
    "        content=user_message\n",
    "    )\n",
    "\n",
    "    # Mappings for partial function calls\n",
    "    call_id_for_index: Dict[int, str] = {}\n",
    "    partial_calls_by_index: Dict[int, dict] = {}\n",
    "    partial_calls_by_id: Dict[str, dict] = {}\n",
    "    in_progress_tools: Dict[str, ChatMessage] = {}\n",
    "\n",
    "    # Titles for tool bubbles\n",
    "    function_titles = {\n",
    "        \"fetch_weather\": \"☁️ fetching weather\",\n",
    "        \"fetch_datetime\": \"🕒 fetching datetime\",\n",
    "        \"fetch_stock_price\": \"📈 fetching financial info\",\n",
    "        \"send_email\": \"✉️ sending mail\",\n",
    "        \"file_search\": \"📄 searching docs\",\n",
    "        \"bing_grounding\": \"🔍 searching bing\",\n",
    "    }\n",
    "\n",
    "    def get_function_title(fn_name: str) -> str:\n",
    "        return function_titles.get(fn_name, f\"🛠 calling {fn_name}\")\n",
    "\n",
    "    def accumulate_args(storage: dict, name_chunk: str, arg_chunk: str):\n",
    "        \"\"\"Accumulates partial JSON data for a function call.\"\"\"\n",
    "        if name_chunk:\n",
    "            storage[\"name\"] += name_chunk\n",
    "        if arg_chunk:\n",
    "            storage[\"args\"] += arg_chunk\n",
    "\n",
    "    def finalize_tool_call(call_id: str):\n",
    "        \"\"\"Creates or updates the ChatMessage bubble for a function call.\"\"\"\n",
    "        if call_id not in partial_calls_by_id:\n",
    "            return\n",
    "        data = partial_calls_by_id[call_id]\n",
    "        fn_name = data[\"name\"].strip()\n",
    "        fn_args = data[\"args\"].strip()\n",
    "        if not fn_name:\n",
    "            return\n",
    "\n",
    "        if call_id not in in_progress_tools:\n",
    "            # Create a new bubble with status=\"pending\"\n",
    "            msg_obj = ChatMessage(\n",
    "                role=\"assistant\",\n",
    "                content=fn_args or \"\",\n",
    "                metadata={\n",
    "                    \"title\": get_function_title(fn_name),\n",
    "                    \"status\": \"pending\",\n",
    "                    \"id\": f\"tool-{call_id}\"\n",
    "                }\n",
    "            )\n",
    "            conversation.append(msg_obj)\n",
    "            in_progress_tools[call_id] = msg_obj\n",
    "        else:\n",
    "            # Update existing bubble\n",
    "            msg_obj = in_progress_tools[call_id]\n",
    "            msg_obj.content = fn_args or \"\"\n",
    "            msg_obj.metadata[\"title\"] = get_function_title(fn_name)\n",
    "\n",
    "    def upsert_tool_call(tcall: dict):\n",
    "        \"\"\"\n",
    "        1) Check the call type\n",
    "        2) If \"function\", gather partial name/args\n",
    "        3) If \"bing_grounding\" or \"file_search\", show a pending bubble\n",
    "        \"\"\"\n",
    "        t_type = tcall.get(\"type\", \"\")\n",
    "        call_id = tcall.get(\"id\")\n",
    "\n",
    "        # --- BING GROUNDING ---\n",
    "        if t_type == \"bing_grounding\":\n",
    "            request_url = tcall.get(\"bing_grounding\", {}).get(\"requesturl\", \"\")\n",
    "            if not request_url.strip():\n",
    "                return\n",
    "\n",
    "            query_str = extract_bing_query(request_url)\n",
    "            if not query_str.strip():\n",
    "                return\n",
    "\n",
    "            msg_obj = ChatMessage(\n",
    "                role=\"assistant\",\n",
    "                content=query_str,\n",
    "                metadata={\n",
    "                    \"title\": get_function_title(\"bing_grounding\"),\n",
    "                    \"status\": \"pending\",\n",
    "                    \"id\": f\"tool-{call_id}\" if call_id else \"tool-noid\"\n",
    "                }\n",
    "            )\n",
    "            conversation.append(msg_obj)\n",
    "            if call_id:\n",
    "                in_progress_tools[call_id] = msg_obj\n",
    "            return\n",
    "\n",
    "        # --- FILE SEARCH ---\n",
    "        elif t_type == \"file_search\":\n",
    "            msg_obj = ChatMessage(\n",
    "                role=\"assistant\",\n",
    "                content=\"searching docs...\",\n",
    "                metadata={\n",
    "                    \"title\": get_function_title(\"file_search\"),\n",
    "                    \"status\": \"pending\",\n",
    "                    \"id\": f\"tool-{call_id}\" if call_id else \"tool-noid\"\n",
    "                }\n",
    "            )\n",
    "            conversation.append(msg_obj)\n",
    "            if call_id:\n",
    "                in_progress_tools[call_id] = msg_obj\n",
    "            return\n",
    "\n",
    "        # --- NON-FUNCTION CALLS ---\n",
    "        elif t_type != \"function\":\n",
    "            return\n",
    "\n",
    "        # --- FUNCTION CALL PARTIAL-ARGS ---\n",
    "        index = tcall.get(\"index\")\n",
    "        new_call_id = call_id\n",
    "        fn_data = tcall.get(\"function\", {})\n",
    "        name_chunk = fn_data.get(\"name\", \"\")\n",
    "        arg_chunk = fn_data.get(\"arguments\", \"\")\n",
    "\n",
    "        if new_call_id:\n",
    "            call_id_for_index[index] = new_call_id\n",
    "\n",
    "        call_id = call_id_for_index.get(index)\n",
    "        if not call_id:\n",
    "            # Accumulate partial\n",
    "            if index not in partial_calls_by_index:\n",
    "                partial_calls_by_index[index] = {\"name\": \"\", \"args\": \"\"}\n",
    "            accumulate_args(partial_calls_by_index[index], name_chunk, arg_chunk)\n",
    "            return\n",
    "\n",
    "        if call_id not in partial_calls_by_id:\n",
    "            partial_calls_by_id[call_id] = {\"name\": \"\", \"args\": \"\"}\n",
    "\n",
    "        if index in partial_calls_by_index:\n",
    "            old_data = partial_calls_by_index.pop(index)\n",
    "            partial_calls_by_id[call_id][\"name\"] += old_data.get(\"name\", \"\")\n",
    "            partial_calls_by_id[call_id][\"args\"] += old_data.get(\"args\", \"\")\n",
    "\n",
    "        # Accumulate partial\n",
    "        accumulate_args(partial_calls_by_id[call_id], name_chunk, arg_chunk)\n",
    "\n",
    "        # Create/update the function bubble\n",
    "        finalize_tool_call(call_id)\n",
    "\n",
    "    # -- EVENT STREAMING --\n",
    "    with project_client.agents.create_stream(\n",
    "        thread_id=thread.id,\n",
    "        assistant_id=agent.id,\n",
    "        event_handler=MyEventHandler()  # the event handler handles console output\n",
    "    ) as stream:\n",
    "        for item in stream:\n",
    "            event_type, event_data, *_ = item\n",
    "\n",
    "            # Remove any None items that might have been appended\n",
    "            conversation = [m for m in conversation if m is not None]\n",
    "\n",
    "            # 1) Partial tool calls\n",
    "            if event_type == \"thread.run.step.delta\":\n",
    "                step_delta = event_data.get(\"delta\", {}).get(\"step_details\", {})\n",
    "                if step_delta.get(\"type\") == \"tool_calls\":\n",
    "                    for tcall in step_delta.get(\"tool_calls\", []):\n",
    "                        upsert_tool_call(tcall)\n",
    "                    yield conversation, \"\"\n",
    "\n",
    "            # 2) run_step\n",
    "            elif event_type == \"run_step\":\n",
    "                step_type = event_data[\"type\"]\n",
    "                step_status = event_data[\"status\"]\n",
    "\n",
    "                # If tool calls are in progress, new or partial\n",
    "                if step_type == \"tool_calls\" and step_status == \"in_progress\":\n",
    "                    for tcall in event_data[\"step_details\"].get(\"tool_calls\", []):\n",
    "                        upsert_tool_call(tcall)\n",
    "                    yield conversation, \"\"\n",
    "\n",
    "                elif step_type == \"tool_calls\" and step_status == \"completed\":\n",
    "                    for cid, msg_obj in in_progress_tools.items():\n",
    "                        msg_obj.metadata[\"status\"] = \"done\"\n",
    "                    in_progress_tools.clear()\n",
    "                    partial_calls_by_id.clear()\n",
    "                    partial_calls_by_index.clear()\n",
    "                    call_id_for_index.clear()\n",
    "                    yield conversation, \"\"\n",
    "\n",
    "                elif step_type == \"message_creation\" and step_status == \"in_progress\":\n",
    "                    msg_id = event_data[\"step_details\"][\"message_creation\"].get(\"message_id\")\n",
    "                    if msg_id:\n",
    "                        conversation.append(ChatMessage(role=\"assistant\", content=\"\"))\n",
    "                    yield conversation, \"\"\n",
    "\n",
    "                elif step_type == \"message_creation\" and step_status == \"completed\":\n",
    "                    yield conversation, \"\"\n",
    "\n",
    "            # 3) partial text from the assistant\n",
    "            elif event_type == \"thread.message.delta\":\n",
    "                agent_msg = \"\"\n",
    "                for chunk in event_data[\"delta\"][\"content\"]:\n",
    "                    agent_msg += chunk[\"text\"].get(\"value\", \"\")\n",
    "\n",
    "                message_id = event_data[\"id\"]\n",
    "\n",
    "                # Try to find a matching assistant bubble\n",
    "                matching_msg = None\n",
    "                for msg in reversed(conversation):\n",
    "                    if msg.metadata and msg.metadata.get(\"id\") == message_id and msg.role == \"assistant\":\n",
    "                        matching_msg = msg\n",
    "                        break\n",
    "\n",
    "                if matching_msg:\n",
    "                    # Append newly streamed text\n",
    "                    matching_msg.content += agent_msg\n",
    "                else:\n",
    "                    # Append to last assistant or create new\n",
    "                    if (\n",
    "                        not conversation\n",
    "                        or conversation[-1].role != \"assistant\"\n",
    "                        or (\n",
    "                            conversation[-1].metadata\n",
    "                            and str(conversation[-1].metadata.get(\"id\", \"\")).startswith(\"tool-\")\n",
    "                        )\n",
    "                    ):\n",
    "                        conversation.append(ChatMessage(role=\"assistant\", content=agent_msg))\n",
    "                    else:\n",
    "                        conversation[-1].content += agent_msg\n",
    "\n",
    "                yield conversation, \"\"\n",
    "\n",
    "            # 4) If entire assistant message is completed\n",
    "            elif event_type == \"thread.message\":\n",
    "                if event_data[\"role\"] == \"assistant\" and event_data[\"status\"] == \"completed\":\n",
    "                    for cid, msg_obj in in_progress_tools.items():\n",
    "                        msg_obj.metadata[\"status\"] = \"done\"\n",
    "                    in_progress_tools.clear()\n",
    "                    partial_calls_by_id.clear()\n",
    "                    partial_calls_by_index.clear()\n",
    "                    call_id_for_index.clear()\n",
    "                    yield conversation, \"\"\n",
    "\n",
    "            # 5) Final done\n",
    "            elif event_type == \"thread.message.completed\":\n",
    "                for cid, msg_obj in in_progress_tools.items():\n",
    "                    msg_obj.metadata[\"status\"] = \"done\"\n",
    "                in_progress_tools.clear()\n",
    "                partial_calls_by_id.clear()\n",
    "                partial_calls_by_index.clear()\n",
    "                call_id_for_index.clear()\n",
    "                yield conversation, \"\"\n",
    "                break\n",
    "\n",
    "    return conversation, \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build a Gradio UI\n",
    "Create a Gradio interface for interacting with the enterprise agent. \n",
    "Include a chatbot component and a text input box for user queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "brand_theme = gr.themes.Default(\n",
    "    primary_hue=\"blue\",\n",
    "    secondary_hue=\"blue\",\n",
    "    neutral_hue=\"gray\",\n",
    "    font=[\"Segoe UI\", \"Arial\", \"sans-serif\"],\n",
    "    font_mono=[\"Courier New\", \"monospace\"],\n",
    "    text_size=\"lg\",\n",
    ").set(\n",
    "    button_primary_background_fill=\"#0f6cbd\",\n",
    "    button_primary_background_fill_hover=\"#115ea3\",\n",
    "    button_primary_background_fill_hover_dark=\"#4f52b2\",\n",
    "    button_primary_background_fill_dark=\"#5b5fc7\",\n",
    "    button_primary_text_color=\"#ffffff\",\n",
    "    button_secondary_background_fill=\"#e0e0e0\",\n",
    "    button_secondary_background_fill_hover=\"#c0c0c0\",\n",
    "    button_secondary_background_fill_hover_dark=\"#a0a0a0\",\n",
    "    button_secondary_text_color=\"#000000\",\n",
    "    body_background_fill=\"#f5f5f5\",\n",
    "    block_background_fill=\"#ffffff\",\n",
    "    body_text_color=\"#242424\",\n",
    "    body_text_color_subdued=\"#616161\",\n",
    "    block_border_color=\"#d1d1d1\",\n",
    "    block_border_color_dark=\"#333333\",\n",
    "    input_background_fill=\"#ffffff\",\n",
    "    input_border_color=\"#d1d1d1\",\n",
    "    input_border_color_focus=\"#0f6cbd\",\n",
    ")\n",
    "\n",
    "with gr.Blocks(theme=brand_theme, css=\"footer {visibility: hidden;}\", fill_height=True) as demo:\n",
    "\n",
    "    def clear_thread():\n",
    "        global thread\n",
    "        thread = project_client.agents.create_thread()\n",
    "        return []\n",
    "\n",
    "    def on_example_clicked(evt: gr.SelectData):\n",
    "        return evt.value[\"text\"]  # Fill the textbox with that example text\n",
    "\n",
    "    gr.HTML(\"<h1 style=\\\"text-align: center;\\\">Azure AI Agent Service</h1>\")\n",
    "\n",
    "    chatbot = gr.Chatbot(\n",
    "        type=\"messages\",\n",
    "        examples=[\n",
    "            {\"text\": \"What's my company's remote work policy?\"},\n",
    "            {\"text\": \"Check if it will rain tomorrow?\"},\n",
    "            {\"text\": \"How is Contoso's stock doing today?\"},\n",
    "            {\"text\": \"Send my direct report a summary of the HR policy.\"},\n",
    "        ],\n",
    "        show_label=False,\n",
    "        scale=1,\n",
    "    )\n",
    "\n",
    "    textbox = gr.Textbox(\n",
    "        show_label=False,\n",
    "        lines=1,\n",
    "        submit_btn=True,\n",
    "    )\n",
    "\n",
    "    # Populate textbox when an example is clicked\n",
    "    chatbot.example_select(fn=on_example_clicked, inputs=None, outputs=textbox)\n",
    "\n",
    "    # On submit: call azure_enterprise_chat, then clear the textbox\n",
    "    (textbox\n",
    "     .submit(\n",
    "         fn=azure_enterprise_chat,\n",
    "         inputs=[textbox, chatbot],\n",
    "         outputs=[chatbot, textbox],\n",
    "     )\n",
    "     .then(\n",
    "         fn=lambda: \"\",\n",
    "         outputs=textbox,\n",
    "     )\n",
    "    )\n",
    "\n",
    "    # A \"Clear\" button that resets the thread and the Chatbot\n",
    "    chatbot.clear(fn=clear_thread, outputs=chatbot)\n",
    "\n",
    "# Launch your Gradio app\n",
    "if __name__ == \"__main__\":\n",
    "    demo.launch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Optional) delete agent, thread, and vector store resources\n",
    "Uncomment out the next cell block to delete the resources created in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#from azure.identity import DefaultAzureCredential\n",
    "#from azure.ai.projects import AIProjectClient\n",
    "#import os\n",
    "#\n",
    "#credential = DefaultAzureCredential()\n",
    "#project_client_delete = AIProjectClient.from_connection_string(\n",
    "#    credential=credential,\n",
    "#    conn_str=os.environ.get(\"PROJECT_CONNECTION_STRING\")\n",
    "#)\n",
    "#\n",
    "#try:\n",
    "#    project_client_delete.agents.delete_agent(agent.id)\n",
    "#    print(\"Agent deletion successful.\")\n",
    "#    project_client_delete.agents.delete_thread(thread.id)\n",
    "#    print(\"Thread deletion successful.\")\n",
    "#    project_client_delete.agents.delete_vector_store(vector_store_id)\n",
    "#    print(\"Vector store deletion successful.\")\n",
    "#    print(\"All deletions succeeded.\")\n",
    "#except Exception as e:\n",
    "#    print(f\"Error during deletion: {e}\")"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}