add support for azure managed identity auth

docs/tutorials/quickstart_azure.ipynb

@@ -0,0 +1,268 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": [
+     "remove-cell"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "# Load from parent directory if not installed\n",
+    "import importlib\n",
+    "import os\n",
+    "\n",
+    "if not importlib.util.find_spec(\"sammo\"):\n",
+    "    import sys\n",
+    "\n",
+    "    sys.path.append(\"../../\")\n",
+    "os.environ[\"CACHE_FILE\"] = \"cache/quickstart.tsv\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "# 🚀 Quick Start\n",
+    "\n",
+    "At the core of SAMMO are symbolic prompt programs. This tutorial will show you a few simple programs.\n",
+    "\n",
+    "To run this example, you need credentials to an AzureOpenAI API compatible model. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# %load -r 3:25 _init.py\n",
+    "import pathlib\n",
+    "import sammo\n",
+    "from sammo.runners import AzureChat\n",
+    "from sammo.base import Template, EvaluationScore\n",
+    "from sammo.components import Output, GenerateText, ForEach, Union\n",
+    "from sammo.extractors import ExtractRegex\n",
+    "from sammo.data import DataTable\n",
+    "import json\n",
+    "import requests\n",
+    "import os\n",
+    "from azure.identity import DefaultAzureCredential, get_bearer_token_provider  \n",
+    "\n",
+    "if not 'AZURE_OPENAI_ENDPOINT' in os.environ:\n",
+    "    raise ValueError(\"Please set the environment variable 'AZURE_OPENAI_ENDPOINT'.\")\n",
+    "\n",
+    "if not 'AZURE_OPENAI_ENDPOINT_DEPLOYMENT' in os.environ:\n",
+    "    raise ValueError(\"Please set the environment variable 'AZURE_OPENAI_ENDPOINT_DEPLOYMENT'.\")\n",
+    "     \n",
+    "# Initialize Azure OpenAI Service client with Entra ID authentication\n",
+    "# similar as https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity#chat-completions\n",
+    "token_provider = get_bearer_token_provider(  \n",
+    "    DefaultAzureCredential(),  \n",
+    "    \"https://cognitiveservices.azure.com/.default\"  \n",
+    ")  \n",
+    "  \n",
+    "api_config = {\n",
+    "    \"azure_ad_token_provider\": token_provider,\n",
+    "    \"endpoint\": os.environ[\"AZURE_OPENAI_ENDPOINT\"],\n",
+    "    \"api_version\": \"2024-05-01-preview\",  # set to whatever version needed\n",
+    "    \"deployment_id\": os.environ[\"AZURE_OPENAI_ENDPOINT_DEPLOYMENT\"],\n",
+    "}\n",
+    "\n",
+    "_ = sammo.setup_logger(\"WARNING\")  # we're only interested in warnings for now\n",
+    "\n",
+    "runner = AzureChat(\n",
+    "    model_id=api_config[\"deployment_id\"],\n",
+    "    api_config=api_config,\n",
+    "    cache=os.getenv(\"CACHE_FILE\", \"cache.tsv\"),\n",
+    "    timeout=30,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "Let's write our first symbolic prompt program (SPP)! How about a quick 'Hello World?'?"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "spp_hello_world = Output(GenerateText(\"Hello World!\"))\n",
+    "spp_hello_world.run(runner)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "## Writing symbolic prompt programs\n",
+    "Okay, let's move on to a more interesting example. For a list of countries, we want the top reason to visit:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "COUNTRIES = [\"Switzerland\", \"Morocco\", \"Tanzania\", \"Indonesia\", \"Peru\"]\n",
+    "\n",
+    "reason_to_visit = GenerateText(\n",
+    "    Template(\"What is the top reason to visit {{input}} in one sentence?\")\n",
+    ")\n",
+    "Output(reason_to_visit).run(runner, COUNTRIES)[:2]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "What happens under the hood is that SAMMO parallizes the execution across all inputs automatically! \n",
+    "\n",
+    "Let's add the best time to visit to it and combine both pieces of information."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "when_to_visit = GenerateText(\n",
+    "    Template(\n",
+    "        \"Which season is the best time to visit {{input}}? Answer in one sentence.\"\n",
+    "    )\n",
+    ")\n",
+    "country_pages = Template(\n",
+    "    \"# {{input}}\\n{{reason}}\\n\\n## When to Visit\\n{{when}}\",\n",
+    "    reason=reason_to_visit,\n",
+    "    when=when_to_visit,\n",
+    ")\n",
+    "written_pages = Output(country_pages).run(runner, COUNTRIES)\n",
+    "written_pages[:2]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Nice! To see what operations SAMMO executed under the hood, we can use `plot_call_trace()` to show all intermediate calls. You can click on each node to get more info."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "written_pages.outputs[0].plot_call_trace()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "## Recap\n",
+    "Let's talk about some of the key concepts from SAMMO we have used:\n",
+    "\n",
+    "1. We constructed a **symbolic prompt program** — a dynamic prompt that is re-used for different inputs.\n",
+    "2. This SPP has a structure which was constructed by nesting **components** from SAMMO.\n",
+    "3. To get the **output**, we call `run()` on the Output component which returns a DataTable.\n",
+    "4. SAMMO **parallelized** execution for us on the input data — no extra work was needed! "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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"
+  },
+  "widgets": {
+   "application/vnd.jupyter.widget-state+json": {
+    "state": {},
+    "version_major": 2,
+    "version_minor": 0
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

pyproject.toml

@@ -59,6 +59,12 @@ sphinx-autodoc2 = "*"
 jupyterlab = "^4.0"
 poetry-core = "^1.8.1"  # (see set_version.py)
 
+[tool.poetry.group.azure]
+optional = true
+
+[tool.poetry.group.azure.dependencies]
+azure-identity = "^1.20"
+
 [tool.black]
 line-length = 120
 

sammo/runners.py

@@ -480,7 +480,13 @@ def _post_init(self):
             self._api_config["endpoint"] = self._api_config["endpoint"][:-1]
 
     def _get_headers(self):
-        return {"api-key": self._api_config["api_key"], "Content-Type": "application/json"}
+        headers = {"Content-Type": "application/json"}
+        if "api_key" in self._api_config:
+            headers = {"api_key": self._api_config["api_key"]}
+        elif "azure_ad_token_provider":
+            # https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity#chat-completions
+            headers = {"Authorization": f"Bearer {self._api_config["azure_ad_token_provider"]()}"}
+        return headers
 
     def _rest_url(self):
         return (