run-llama--llama_index/docs/examples/vector_stores/azurepostgresql.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2dd8cd46",
   "metadata": {},
   "source": [
    "# Azure Postgres Vector Store\n",
    "In this notebook we are going to show how to use [Azure Postgresql](https://azure.microsoft.com/en-au/products/postgresql) and  [pg_diskann](https://github.com/microsoft/DiskANN) to perform vector searches in LlamaIndex. \n",
    "Please note that this document is mostly based on the document for [PostgreSQL integration](https://docs.llamaindex.ai/en/stable/examples/vector_stores/postgres/) to simplify the transition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d4b9721",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c95fd172",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext sql"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3412ab2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "import os\n",
    "from urllib.parse import quote_plus\n",
    "\n",
    "cmd = [\n",
    "    \"az\",\n",
    "    \"account\",\n",
    "    \"get-access-token\",\n",
    "    \"--resource\",\n",
    "    \"https://ossrdbms-aad.database.windows.net\",\n",
    "    \"--query\",\n",
    "    \"accessToken\",\n",
    "    \"--output\",\n",
    "    \"tsv\",\n",
    "]\n",
    "\n",
    "try:\n",
    "    token = subprocess.check_output(cmd, text=True).strip()\n",
    "except subprocess.CalledProcessError as exc:\n",
    "    raise RuntimeError(f\"Failed to run command: {exc}\") from exc\n",
    "os.environ[\"PGPASSWORD\"] = token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa5389a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<span style=\"None\">Connecting to &#x27;postgresql://&#x27;</span>"
      ],
      "text/plain": [
       "Connecting to 'postgresql://'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%sql postgresql://"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a35ff87e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<span style=\"None\">Running query in &#x27;postgresql://&#x27;</span>"
      ],
      "text/plain": [
       "Running query in 'postgresql://'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <thead>\n",
       "        <tr>\n",
       "        </tr>\n",
       "    </thead>\n",
       "    <tbody>\n",
       "    </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "++\n",
       "||\n",
       "++\n",
       "++"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%sql\n",
    "drop table if exists llamaindex_vectors;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ece76c79",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "import os\n",
    "\n",
    "# Uncomment to see debug logs\n",
    "# logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)\n",
    "# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index.core import (\n",
    "    SimpleDirectoryReader,\n",
    "    StorageContext,\n",
    "    VectorStoreIndex,\n",
    ")\n",
    "from llama_index.core.settings import Settings\n",
    "from llama_index.llms.azure_openai import AzureOpenAI\n",
    "from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding\n",
    "import textwrap\n",
    "\n",
    "# Import from the local file\n",
    "from llama_index.vector_stores.azure_postgres import AzurePGVectorStore\n",
    "from llama_index.vector_stores.azure_postgres.common import (\n",
    "    AzurePGConnectionPool,\n",
    "    DiskANN,\n",
    "    VectorOpClass,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdcc750f",
   "metadata": {},
   "source": [
    "### Setup OpenAI\n",
    "The first step is to configure the Azure openai key. It will be used to created embeddings for the documents loaded into the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2517eb0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Method 1: Using os.environ.get() with fallback values\n",
    "aoai_api_key = os.environ.get(\"AOAI_API_KEY\", \"key\")\n",
    "aoai_endpoint = os.environ.get(\"AOAI_ENDPOINT\", \"endpoint\")\n",
    "aoai_api_version = os.environ.get(\"AOAI_API_VERSION\", \"2024-12-01-preview\")\n",
    "\n",
    "llm = AzureOpenAI(\n",
    "    model=\"o4-mini\",\n",
    "    deployment_name=\"o4-mini\",\n",
    "    api_key=aoai_api_key,\n",
    "    azure_endpoint=aoai_endpoint,\n",
    "    api_version=aoai_api_version,\n",
    ")\n",
    "\n",
    "# You need to deploy your own embedding model as well as your own chat completion model\n",
    "embed_model = AzureOpenAIEmbedding(\n",
    "    model=\"text-embedding-3-small\",\n",
    "    deployment_name=\"text-embedding-3-small\",\n",
    "    api_key=aoai_api_key,\n",
    "    azure_endpoint=aoai_endpoint,\n",
    "    api_version=aoai_api_version,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ef64e3f",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9644d4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2025-09-03 15:56:56--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.111.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n",
      "\n",
      "data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.1s    \n",
      "\n",
      "2025-09-03 15:56:56 (765 KB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72c1056a",
   "metadata": {},
   "source": [
    "### Loading documents\n",
    "Load the documents stored in the `data/paul_graham/` using the SimpleDirectoryReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7e41e89",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document ID: 4a7a27c2-6013-408b-aa3d-65fd89b824d8\n"
     ]
    }
   ],
   "source": [
    "documents = SimpleDirectoryReader(\"./data/paul_graham\").load_data()\n",
    "print(\"Document ID:\", documents[0].doc_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4aa5bc3",
   "metadata": {},
   "source": [
    "### Create the Database\n",
    "Using an existing postgres instance running on Azure, we will use Microsoft Entra authentication to connect to the database. Please make sure you are logged in to your Azure account.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6ce826a",
   "metadata": {},
   "outputs": [],
   "source": [
    "host = os.environ.get(\"PGHOST\", \"<your_host>\")\n",
    "port = int(os.environ.get(\"PGPORT\", 5432))\n",
    "database = os.environ.get(\"PGDATABASE\", \"postgres\")\n",
    "from psycopg import Connection\n",
    "from psycopg.rows import dict_row\n",
    "from llama_index.vector_stores.azure_postgres.common import (\n",
    "    ConnectionInfo,\n",
    "    create_extensions,\n",
    "    Extension,\n",
    ")\n",
    "\n",
    "\n",
    "def configure_connection(conn: Connection) -> None:\n",
    "    conn.autocommit = True\n",
    "    create_extensions(conn, [Extension(ext_name=\"vector\")])\n",
    "    create_extensions(conn, [Extension(ext_name=\"pg_diskann\")])\n",
    "    conn.row_factory = dict_row\n",
    "\n",
    "\n",
    "azure_conn_info: ConnectionInfo = ConnectionInfo(\n",
    "    host=host, port=port, dbname=database, configure=configure_connection\n",
    ")\n",
    "conn = AzurePGConnectionPool(\n",
    "    azure_conn_info=azure_conn_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "374d3e11",
   "metadata": {},
   "source": [
    "### Create the vector store\n",
    "Here we create an index backed by Postgres using the documents loaded previously. AzurePGVectorStore takes a few arguments. The example below constructs a PGVectorStore with no index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a17c8526",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding type is not specified, defaulting to 'vector'.\n",
      "Embedding dimension is not specified, defaulting to 1536.\n",
      "Embedding index is not specified, defaulting to 'DiskANN' with 'vector_cosine_ops' opclass.\n",
      "/home/kislalorhan/workspace/myenv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Parsing nodes: 100%|██████████| 1/1 [00:00<00:00, 11.88it/s]\n",
      "Generating embeddings: 100%|██████████| 22/22 [00:02<00:00,  9.54it/s]\n"
     ]
    }
   ],
   "source": [
    "vector_store = AzurePGVectorStore.from_params(\n",
    "    connection_pool=conn,\n",
    "    table_name=\"llamaindex_vectors\",\n",
    "    embed_dim=1536,  # openai embedding dimension\n",
    ")\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context, show_progress=True\n",
    ")\n",
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc153851",
   "metadata": {},
   "source": [
    "### Query the dataset\n",
    "We can now ask questions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c7ba58f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "He pursued two parallel creative tracks—writing and programming.   • As a teenager he wrote\n",
      "(admittedly “awful”) short stories and taught himself to program on his school’s IBM 1401, later\n",
      "moving on to a TRS-80 where he wrote simple games, a model-rocket flight predictor, and even a small\n",
      "word processor.   • In college he initially majored in philosophy but switched to AI, became\n",
      "fascinated by Lisp, and decided to write a book on Lisp hacking.  Much of what became On Lisp was\n",
      "drafted during his grad-school years.   • At the same time, seeking a more permanent art form, he\n",
      "began taking painting classes at Harvard, planning to make and earn a living from paintings that,\n",
      "unlike software, wouldn’t become obsolete.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"What did the author do?\")\n",
    "print(textwrap.fill(str(response), 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2211642f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Artificial intelligence became a hot topic. Two specific influences drove that surge of interest: -\n",
      "Heinlein’s science-fiction novel The Moon Is a Harsh Mistress, featuring the self-aware computer\n",
      "“Mike”   - A PBS documentary demonstrating Terry Winograd’s SHRDLU natural-language program\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"What happened in the mid 1980s?\")\n",
    "print(textwrap.fill(str(response), 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5e0d3af",
   "metadata": {},
   "source": [
    "### Querying existing index\n",
    "Now, we create a pg_diskann index with max_neighbors = 32, l_value_ib = 100, and l_value_is = 100, with the `vector_cosine_ops` method on our embeddings and use it with a new vector store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3098561",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<span style=\"None\">Running query in &#x27;postgresql://&#x27;</span>"
      ],
      "text/plain": [
       "Running query in 'postgresql://'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <thead>\n",
       "        <tr>\n",
       "        </tr>\n",
       "    </thead>\n",
       "    <tbody>\n",
       "    </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "++\n",
       "||\n",
       "++\n",
       "++"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%sql \n",
    "create index on llamaindex_vectors \n",
    "using diskann (embedding vector_cosine_ops) \n",
    "with (\n",
    "\tmax_neighbors = 32,\n",
    "\tl_value_ib = 100\n",
    ");\n",
    "set diskann.l_value_is to 100;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00e6490d",
   "metadata": {},
   "outputs": [],
   "source": [
    "diskann = DiskANN(\n",
    "    op_class=VectorOpClass.vector_cosine_ops,\n",
    "    max_neighbors=32,\n",
    "    l_value_ib=100,\n",
    "    l_value_is=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f63d515",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'schema_name': 'public', 'table_name': 'llamaindex_vectors', 'index_name': 'llamaindex_vectors_embedding_idx', 'index_type': 'diskann', 'index_column': 'embedding', 'index_opclass': 'vector_cosine_ops', 'index_opts': ['max_neighbors=32', 'l_value_ib=100']}]\n"
     ]
    }
   ],
   "source": [
    "vector_store = AzurePGVectorStore.from_params(\n",
    "    connection_pool=conn,\n",
    "    schema_name=\"public\",\n",
    "    table_name=\"llamaindex_vectors\",\n",
    "    embed_dim=1536,  # openai embedding dimension\n",
    "    embedding_index=diskann,\n",
    ")\n",
    "\n",
    "index = VectorStoreIndex.from_vector_store(vector_store=vector_store)\n",
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60569a74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "He spent his spare time writing (mostly really bad short stories) and learning to program.  As a\n",
      "teenager he punched out Fortran jobs on an IBM 1401, then moved on to a TRS-80 microcomputer, where\n",
      "he wrote simple games, a model-rocket flight predictor, and even a tiny word-processor.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"What did the author do?\")\n",
    "print(textwrap.fill(str(response), 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3966d7bb",
   "metadata": {},
   "source": [
    "### Access individual nodes\n",
    "Read a specific node by its id."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab3dedd8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22\n",
      "3dd2f695-1def-431b-ae2c-2561472a0272\n",
      "Node ID: 3dd2f695-1def-431b-ae2c-2561472a0272\n",
      "Text: What I Worked On  February 2021  Before college the two main\n",
      "things I worked on, outside of school, were writing and programming. I\n",
      "didn't write essays. I wrote what beginning writers were supposed to\n",
      "write then, and probably still are: short stories. My stories were\n",
      "awful. They had hardly any plot, just characters with strong feelings,\n",
      "which I ...\n"
     ]
    }
   ],
   "source": [
    "nodes = vector_store.get_nodes()\n",
    "print(len(nodes))\n",
    "node_id = nodes[0].node_id\n",
    "print(node_id)\n",
    "nodes = vector_store.get_nodes([node_id])\n",
    "print(nodes[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "695bfe14",
   "metadata": {},
   "source": [
    "Delete a single node and then the whole table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c3e0e23",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "21\n"
     ]
    }
   ],
   "source": [
    "vector_store.delete_nodes(node_ids=[node_id])\n",
    "nodes = vector_store.get_nodes()\n",
    "print(len(nodes))\n",
    "vector_store.clear()  # delete all\n",
    "nodes = vector_store.get_nodes()\n",
    "print(len(nodes))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58c68a15",
   "metadata": {},
   "source": [
    "### Metadata filters\n",
    "\n",
    "AzurePGVectorStore supports storing metadata in nodes, and filtering based on that metadata during the retrieval step."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da4acf74",
   "metadata": {},
   "source": [
    "#### Download git commits dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6368602",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'commit': '03baef1008086ed4960042fa463e570072173bb5', 'author': 'Benjamin Christopher Simmonds <44439583+benibenj@users.noreply.github.com>', 'date': 'Mon Aug 25 13:01:24 2025 +0200', 'change summary': 'Support registering views to the secondary side bar (#261619)', 'change details': \"* Support registering views to the secondary side bar\\\\n\\\\n* rename to secondarySideBar\\\\n\\\\n* Rename 'auxiliarybar' to 'secondarySidebar'\"}\n",
      "169\n"
     ]
    }
   ],
   "source": [
    "# !mkdir -p 'data/csv/'\n",
    "# !wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/csv/commit_history_2.csv' -O 'data/csv/commit_history_2.csv'\n",
    "import builtins\n",
    "import csv\n",
    "\n",
    "# TODO: Once the PR is merged: Change this to with open(\"data/csv/commit_history_2.csv\", \"r\") as f:\n",
    "with builtins.open(\"../data/csv/commit_history_2.csv\", \"r\") as f:\n",
    "    commits = list(csv.DictReader(f))\n",
    "\n",
    "print(commits[0])\n",
    "print(len(commits))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83a1f6fd",
   "metadata": {},
   "source": [
    "#### Add nodes with custom metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "174c093b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Node ID: 9a06a469-32bc-4dd3-90a7-d6b933e7ad3f\n",
      "Text: Support registering views to the secondary side bar (#261619)  *\n",
      "Support registering views to the secondary side bar\\n\\n* rename to\n",
      "secondarySideBar\\n\\n* Rename 'auxiliarybar' to 'secondarySidebar'\n",
      "2025-08-18 to 2025-08-25\n",
      "{'benjamin.pasero@microsoft.com', 'lramos15@gmail.com', 'matb@microsoft.com', 'mpg@mpg.is', '23246594+joshspicer@users.noreply.github.com', '2644648+TylerLeonhardt@users.noreply.github.com', '62267334+anthonykim1@users.noreply.github.com', '2193314+Tyriar@users.noreply.github.com', '3372902+lszomoru@users.noreply.github.com', 'martinae@microsoft.com', '38270282+alexr00@users.noreply.github.com', 'copeet@microsoft.com', 'rwoll@users.noreply.github.com', 'merogge@microsoft.com', '44439583+benibenj@users.noreply.github.com', '4821+timheuer@users.noreply.github.com', '49699333+dependabot[bot]@users.noreply.github.com', '54879025+justschen@users.noreply.github.com', 'bhavyau@microsoft.com', 'roblourens@gmail.com', 'ethanbovard@hotmail.com', 'hkirschner@microsoft.com', 'hop2deep@gmail.com', '198982749+Copilot@users.noreply.github.com', 'amarlenkyzy@microsoft.com'}\n"
     ]
    }
   ],
   "source": [
    "# Create TextNode for each of the first 100 commits\n",
    "from llama_index.core.schema import TextNode\n",
    "from datetime import datetime\n",
    "import re\n",
    "\n",
    "nodes = []\n",
    "dates = set()\n",
    "authors = set()\n",
    "for commit in commits[:100]:\n",
    "    author_email = commit[\"author\"].split(\"<\")[1][:-1]\n",
    "    commit_date = datetime.strptime(\n",
    "        commit[\"date\"], \"%a %b %d %H:%M:%S %Y %z\"\n",
    "    ).strftime(\"%Y-%m-%d\")\n",
    "    commit_text = commit[\"change summary\"]\n",
    "    if commit[\"change details\"]:\n",
    "        commit_text += \"\\n\\n\" + commit[\"change details\"]\n",
    "    fixes = re.findall(r\"#(\\d+)\", commit_text, re.IGNORECASE)\n",
    "    nodes.append(\n",
    "        TextNode(\n",
    "            text=commit_text,\n",
    "            metadata={\n",
    "                \"commit_date\": commit_date,\n",
    "                \"author\": author_email,\n",
    "                \"fixes\": fixes,\n",
    "            },\n",
    "        )\n",
    "    )\n",
    "    dates.add(commit_date)\n",
    "    authors.add(author_email)\n",
    "\n",
    "print(nodes[0])\n",
    "print(min(dates), \"to\", max(dates))\n",
    "print(authors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d76cbe6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding type is not specified, defaulting to 'vector'.\n",
      "Embedding dimension is not specified, defaulting to 1536.\n",
      "Embedding index is not specified, defaulting to 'DiskANN' with 'vector_cosine_ops' opclass.\n"
     ]
    }
   ],
   "source": [
    "vector_store = AzurePGVectorStore.from_params(\n",
    "    connection_pool=conn,\n",
    "    schema_name=\"public\",\n",
    "    table_name=\"metadata_filter_demo3\",\n",
    "    embed_dim=1536,  # openai embedding dimension\n",
    ")\n",
    "\n",
    "index = VectorStoreIndex.from_vector_store(vector_store=vector_store)\n",
    "index.insert_nodes(nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83297f88",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "He added an opt-in “automation” mode that swaps in custom dialog windows (and a simple file picker) so that modal dialogs can be surfaced and driven under automation.\n"
     ]
    }
   ],
   "source": [
    "print(index.as_query_engine().query(\"How did Leonhardt allow modal?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c32eeb2d",
   "metadata": {},
   "source": [
    "#### Apply metadata filters\n",
    "\n",
    "Now we can filter by commit author or by date when retrieving nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f125dfd2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-22', 'author': 'benjamin.pasero@microsoft.com', 'fixes': ['262878']}\n",
      "{'commit_date': '2025-08-21', 'author': 'matb@microsoft.com', 'fixes': ['262772', '262772']}\n",
      "{'commit_date': '2025-08-21', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'benjamin.pasero@microsoft.com', 'fixes': ['262444', '262417']}\n",
      "{'commit_date': '2025-08-25', 'author': 'benjamin.pasero@microsoft.com', 'fixes': ['263211']}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': ['262472']}\n",
      "{'commit_date': '2025-08-22', 'author': 'benjamin.pasero@microsoft.com', 'fixes': ['5761', '262439']}\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.vector_stores.types import (\n",
    "    MetadataFilter,\n",
    "    MetadataFilters,\n",
    ")\n",
    "\n",
    "filters = MetadataFilters(\n",
    "    filters=[\n",
    "        MetadataFilter(key=\"author\", value=\"matb@microsoft.com\"),\n",
    "        MetadataFilter(key=\"author\", value=\"benjamin.pasero@microsoft.com\"),\n",
    "    ],\n",
    "    condition=\"or\",\n",
    ")\n",
    "\n",
    "retriever = index.as_retriever(\n",
    "    similarity_top_k=10,\n",
    "    filters=filters,\n",
    ")\n",
    "\n",
    "retrieved_nodes = retriever.retrieve(\"What is this software project about?\")\n",
    "\n",
    "for node in retrieved_nodes:\n",
    "    print(node.node.metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81d9ca51",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'commit_date': '2025-08-22', 'author': '2644648+TylerLeonhardt@users.noreply.github.com', 'fixes': ['262984']}\n",
      "{'commit_date': '2025-08-22', 'author': '198982749+Copilot@users.noreply.github.com', 'fixes': ['261705']}\n",
      "{'commit_date': '2025-08-20', 'author': 'bhavyau@microsoft.com', 'fixes': ['262619']}\n",
      "{'commit_date': '2025-08-21', 'author': 'merogge@microsoft.com', 'fixes': ['262732', '252515']}\n",
      "{'commit_date': '2025-08-22', 'author': '54879025+justschen@users.noreply.github.com', 'fixes': ['262975']}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-21', 'author': '198982749+Copilot@users.noreply.github.com', 'fixes': ['262214']}\n",
      "{'commit_date': '2025-08-21', 'author': '2644648+TylerLeonhardt@users.noreply.github.com', 'fixes': ['262510']}\n",
      "{'commit_date': '2025-08-21', 'author': '54879025+justschen@users.noreply.github.com', 'fixes': ['262802']}\n",
      "{'commit_date': '2025-08-22', 'author': '54879025+justschen@users.noreply.github.com', 'fixes': ['262951']}\n"
     ]
    }
   ],
   "source": [
    "filters = MetadataFilters(\n",
    "    filters=[\n",
    "        MetadataFilter(key=\"commit_date\", value=\"2025-08-20\", operator=\">=\"),\n",
    "        MetadataFilter(key=\"commit_date\", value=\"2025-08-25\", operator=\"<=\"),\n",
    "    ],\n",
    "    condition=\"and\",\n",
    ")\n",
    "\n",
    "retriever = index.as_retriever(\n",
    "    similarity_top_k=10,\n",
    "    filters=filters,\n",
    ")\n",
    "\n",
    "retrieved_nodes = retriever.retrieve(\"What is this software project about?\")\n",
    "\n",
    "for node in retrieved_nodes:\n",
    "    print(node.node.metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2abff37",
   "metadata": {},
   "source": [
    "#### Apply nested filters\n",
    "\n",
    "In the above examples, we combined multiple filters using AND or OR. We can also combine multiple sets of filters.\n",
    "\n",
    "e.g. in SQL:\n",
    "```sql\n",
    "WHERE (commit_date >= '2025-08-20' AND commit_date <= '2023-08-25') AND (author = 'matb@microsoft.com' OR author = 'benjamin.pasero@microsoft.com')\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4fdf808",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-22', 'author': 'benjamin.pasero@microsoft.com', 'fixes': ['262878']}\n",
      "{'commit_date': '2025-08-21', 'author': 'matb@microsoft.com', 'fixes': ['262772', '262772']}\n",
      "{'commit_date': '2025-08-21', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'benjamin.pasero@microsoft.com', 'fixes': ['262444', '262417']}\n",
      "{'commit_date': '2025-08-25', 'author': 'benjamin.pasero@microsoft.com', 'fixes': ['263211']}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': ['262472']}\n",
      "{'commit_date': '2025-08-22', 'author': 'benjamin.pasero@microsoft.com', 'fixes': ['5761', '262439']}\n"
     ]
    }
   ],
   "source": [
    "filters = MetadataFilters(\n",
    "    filters=[\n",
    "        MetadataFilters(\n",
    "            filters=[\n",
    "                MetadataFilter(\n",
    "                    key=\"commit_date\", value=\"2025-08-20\", operator=\">=\"\n",
    "                ),\n",
    "                MetadataFilter(\n",
    "                    key=\"commit_date\", value=\"2025-08-25\", operator=\"<=\"\n",
    "                ),\n",
    "            ],\n",
    "            condition=\"and\",\n",
    "        ),\n",
    "        MetadataFilters(\n",
    "            filters=[\n",
    "                MetadataFilter(key=\"author\", value=\"matb@microsoft.com\"),\n",
    "                MetadataFilter(\n",
    "                    key=\"author\", value=\"benjamin.pasero@microsoft.com\"\n",
    "                ),\n",
    "            ],\n",
    "            condition=\"or\",\n",
    "        ),\n",
    "    ],\n",
    "    condition=\"and\",\n",
    ")\n",
    "\n",
    "retriever = index.as_retriever(\n",
    "    similarity_top_k=10,\n",
    "    filters=filters,\n",
    ")\n",
    "\n",
    "retrieved_nodes = retriever.retrieve(\"What is this software project about?\")\n",
    "\n",
    "for node in retrieved_nodes:\n",
    "    print(node.node.metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90ce4d6f",
   "metadata": {},
   "source": [
    "The above can be simplified by using the IN operator. `AzurePGVectorStore` supports `in`, `nin`, and `contains` for comparing an element with a list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7cdc024",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'benjamin.pasero@microsoft.com', 'fixes': ['262444', '262417']}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': ['262472']}\n",
      "{'commit_date': '2025-08-18', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-19', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': ['262508']}\n",
      "{'commit_date': '2025-08-20', 'author': 'matb@microsoft.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-18', 'author': 'matb@microsoft.com', 'fixes': ['262219']}\n"
     ]
    }
   ],
   "source": [
    "filters = MetadataFilters(\n",
    "    filters=[\n",
    "        MetadataFilter(key=\"commit_date\", value=\"2025-08-15\", operator=\">=\"),\n",
    "        MetadataFilter(key=\"commit_date\", value=\"2025-08-20\", operator=\"<=\"),\n",
    "        MetadataFilter(\n",
    "            key=\"author\",\n",
    "            value=[\"matb@microsoft.com\", \"benjamin.pasero@microsoft.com\"],\n",
    "            operator=\"in\",\n",
    "        ),\n",
    "    ],\n",
    "    condition=\"and\",\n",
    ")\n",
    "\n",
    "retriever = index.as_retriever(\n",
    "    similarity_top_k=10,\n",
    "    filters=filters,\n",
    ")\n",
    "\n",
    "retrieved_nodes = retriever.retrieve(\"What is this software project about?\")\n",
    "\n",
    "for node in retrieved_nodes:\n",
    "    print(node.node.metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1535db65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'commit_date': '2025-08-20', 'author': 'bhavyau@microsoft.com', 'fixes': ['262619']}\n",
      "{'commit_date': '2025-08-19', 'author': '3372902+lszomoru@users.noreply.github.com', 'fixes': ['262276']}\n",
      "{'commit_date': '2025-08-19', 'author': '54879025+justschen@users.noreply.github.com', 'fixes': ['262239']}\n",
      "{'commit_date': '2025-08-18', 'author': 'roblourens@gmail.com', 'fixes': ['262222', '260539']}\n",
      "{'commit_date': '2025-08-19', 'author': '2644648+TylerLeonhardt@users.noreply.github.com', 'fixes': ['262417']}\n",
      "{'commit_date': '2025-08-19', 'author': '54879025+justschen@users.noreply.github.com', 'fixes': ['262362']}\n",
      "{'commit_date': '2025-08-18', 'author': '2644648+TylerLeonhardt@users.noreply.github.com', 'fixes': ['262260']}\n",
      "{'commit_date': '2025-08-20', 'author': '2644648+TylerLeonhardt@users.noreply.github.com', 'fixes': ['262564']}\n",
      "{'commit_date': '2025-08-19', 'author': '198982749+Copilot@users.noreply.github.com', 'fixes': []}\n",
      "{'commit_date': '2025-08-19', 'author': '198982749+Copilot@users.noreply.github.com', 'fixes': []}\n"
     ]
    }
   ],
   "source": [
    "# Same thing, with NOT IN\n",
    "filters = MetadataFilters(\n",
    "    filters=[\n",
    "        MetadataFilter(key=\"commit_date\", value=\"2025-08-15\", operator=\">=\"),\n",
    "        MetadataFilter(key=\"commit_date\", value=\"2025-08-20\", operator=\"<=\"),\n",
    "        MetadataFilter(\n",
    "            key=\"author\",\n",
    "            value=[\"matb@microsoft.com\", \"benjamin.pasero@microsoft.com\"],\n",
    "            operator=\"nin\",\n",
    "        ),\n",
    "    ],\n",
    "    condition=\"and\",\n",
    ")\n",
    "\n",
    "retriever = index.as_retriever(\n",
    "    similarity_top_k=10,\n",
    "    filters=filters,\n",
    ")\n",
    "\n",
    "retrieved_nodes = retriever.retrieve(\"What is this software project about?\")\n",
    "\n",
    "for node in retrieved_nodes:\n",
    "    print(node.node.metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "220d7939",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CONTAINS\n",
    "filters = MetadataFilters(\n",
    "    filters=[\n",
    "        MetadataFilter(key=\"fixes\", value=\"5680\", operator=\"contains\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "retriever = index.as_retriever(\n",
    "    similarity_top_k=10,\n",
    "    filters=filters,\n",
    ")\n",
    "\n",
    "retrieved_nodes = retriever.retrieve(\"How did these commits fix the issue?\")\n",
    "for node in retrieved_nodes:\n",
    "    print(node.node.metadata)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myenv",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}