{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "x5Y94EG0V_kp" }, "source": [ "# Extract Metadata Filters from a Query\n", "\n", "*Notebook by [David Batista](https://www.linkedin.com/in/dsbatista)*\n", "\n", "> This is part one of the **Advanced Use Cases** series:\n", ">\n", "> 1️⃣ **Extract Metadata from Queries to Improve Retrieval & the [full article](/blog/extracting-metadata-filter)**\n", ">\n", "> 2️⃣ Query Expansion [cookbook](/cookbook/query-expansion) & [full article](/blog/query-expansion)\n", ">\n", "> 3️⃣ Query Decomposition [cookbook](/cookbook/query_decomposition) & the [full article](/blog/query-decomposition)\n", ">\n", "> 4️⃣ [Automated Metadata Enrichment](/cookbook/metadata_enrichment)\n", "\n", "In this notebook, we'll discuss how to implement a custom component, `QueryMetadataExtractor`, that extracts entities from the query and formulates the corresponding metadata filter.\n", "\n", "**Useful Sources**\n", "\n", "* [📖 Docs](https://docs.haystack.deepset.ai/docs/intro)\n", "* [📚 Tutorials](https://haystack.deepset.ai/tutorials)\n", "* [🧑‍🍳 Cookbooks](https://github.com/deepset-ai/haystack-cookbook)" ] }, { "cell_type": "markdown", "metadata": { "id": "UFOEMHACYVJf" }, "source": [ "## Setup the Development Environment" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "hX7LembzP_pU", "outputId": "8c022487-0b00-454e-c825-1eb1faa0b3a7" }, "outputs": [], "source": [ "!pip install haystack-ai\n", "!pip install sentence-transformers" ] }, { "cell_type": "markdown", "metadata": { "id": "SiVPzKKT-DYk" }, "source": [ "Enter your `OPENAI_API_KEY`. Get your OpenAI API key [here](https://platform.openai.com/api-keys):" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "mR9MmczidDzI", "outputId": "43e8a01c-9e7a-4de0-9f94-5ae7a02d0735" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Enter OpenAI API key:··········\n" ] } ], "source": [ "import os\n", "from getpass import getpass\n", "\n", "if \"OPENAI_API_KEY\" not in os.environ:\n", " os.environ[\"OPENAI_API_KEY\"] = getpass(\"Enter OpenAI API key:\")" ] }, { "cell_type": "markdown", "metadata": { "id": "cjS5CnCs-hcA" }, "source": [ "## Implement `QueryMetadataExtractor`\n", "\n", "Create a [custom component](https://docs.haystack.deepset.ai/docs/custom-components), `QueryMetadataExtractor`, which takes `query` and `metadata_fields` as inputs and outputs `filters`. This component encapsulates a generative pipeline, made up of [`PromptBuilder`](https://docs.haystack.deepset.ai/docs/promptbuilder) and [`OpenAIGenerator`](https://docs.haystack.deepset.ai/docs/openaigenerator). The pipeline instructs the LLM to extract keywords, phrases, or entities from a given query which can then be used as metadata filters. In the prompt, we include instructions to ensure the output format is in JSON and provide `metadata_fields` along with the `query` to ensure the correct entities are extracted from the query.\n", "\n", "Once the pipeline is initialized in the `init` method of the component, we post-process the LLM output in the `run` method. This step ensures the extracted metadata is correctly formatted to be used as a metadata filter." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "5eW_YagJUUh9" }, "outputs": [], "source": [ "import json\n", "from typing import Dict, List\n", "\n", "from haystack import Pipeline, component\n", "from haystack.components.builders import PromptBuilder\n", "from haystack.components.generators import OpenAIGenerator\n", "\n", "@component()\n", "class QueryMetadataExtractor:\n", "\n", " def __init__(self):\n", " prompt = \"\"\"\n", " You are part of an information system that processes users queries.\n", " Given a user query you extract information from it that matches a given list of metadata fields.\n", " The information to be extracted from the query must match the semantics associated with the given metadata fields.\n", " The information that you extracted from the query will then be used as filters to narrow down the search space\n", " when querying an index.\n", " Just include the value of the extracted metadata without including the name of the metadata field.\n", " The extracted information in 'Extracted metadata' must be returned as a valid JSON structure.\n", " ###\n", " Example 1:\n", " Query: \"What was the revenue of Nvidia in 2022?\"\n", " Metadata fields: {\"company\", \"year\"}\n", " Extracted metadata fields: {\"company\": \"nvidia\", \"year\": 2022}\n", " ###\n", " Example 2:\n", " Query: \"What were the most influential publications in 2023 regarding Alzheimer's disease?\"\n", " Metadata fields: {\"disease\", \"year\"}\n", " Extracted metadata fields: {\"disease\": \"Alzheimer\", \"year\": 2023}\n", " ###\n", " Example 3:\n", " Query: \"{{query}}\"\n", " Metadata fields: \"{{metadata_fields}}\"\n", " Extracted metadata fields:\n", " \"\"\"\n", " self.pipeline = Pipeline()\n", " self.pipeline.add_component(name=\"builder\", instance=PromptBuilder(prompt))\n", " self.pipeline.add_component(name=\"llm\", instance=OpenAIGenerator(model=\"gpt-4o-mini\"))\n", " self.pipeline.connect(\"builder\", \"llm\")\n", "\n", " @component.output_types(filters=Dict[str, str])\n", " def run(self, query: str, metadata_fields: List[str]):\n", " result = self.pipeline.run({'builder': {'query': query, 'metadata_fields': metadata_fields}})\n", " metadata = json.loads(result['llm']['replies'][0])\n", "\n", " # this can be done with specific data structures and in a more sophisticated way\n", " filters = []\n", " for key, value in metadata.items():\n", " field = f\"meta.{key}\"\n", " filters.append({f\"field\": field, \"operator\": \"==\", \"value\": value})\n", "\n", " return {\"filters\": {\"operator\": \"AND\", \"conditions\": filters}}" ] }, { "cell_type": "markdown", "metadata": { "id": "OFWt-KH__bD-" }, "source": [ "First, let's test the `QueryMetadataExtractor` in isolation, passing a query and a list of metadata fields." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "MvM11z63bOoq", "outputId": "f7f2c44b-136e-45e5-8b64-b96b00962477" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "{'filters': {'operator': 'AND', 'conditions': [{'field': 'meta.year', 'operator': '==', 'value': 2022}, {'field': 'meta.disease', 'operator': '==', 'value': 'Parkinson'}]}}\n" ] } ], "source": [ "extractor = QueryMetadataExtractor()\n", "\n", "query = \"What were the most influential publications in 2022 regarding Parkinson's disease?\"\n", "metadata_fields = {\"disease\", \"year\"}\n", "\n", "result = extractor.run(query, metadata_fields)\n", "print(result)" ] }, { "cell_type": "markdown", "metadata": { "id": "wAw4QeMN_dfM" }, "source": [ "Notice that the `QueryMetadataExtractor` has extracted the metadata fields from the query and returned them in a format that can be used as filters passed directly to a `Retriever`. By default, the `QueryMetadataExtractor` will use all metadata fields as conditions together with an `AND` operator." ] }, { "cell_type": "markdown", "metadata": { "id": "wmibB-k3_v5D" }, "source": [ "## Use `QueryMetadataExtractor` in a Pipeline\n", "\n", "Now, let's plug the `QueryMetadataExtractor` into a `Pipeline` with a `Retriever` connected to a `DocumentStore` to see how it works in practice.\n", "\n", "We start by creating a `InMemoryDocumentStore` and adding some documents to it. We include info about “year” and “disease” in the “meta” field of each document." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "ICoV0EbII3KK", "outputId": "d648e721-a293-48fa-b5ed-307ef39f70a5" }, "outputs": [ { "data": { "text/plain": [ "4" ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from haystack import Document\n", "from haystack.document_stores.in_memory import InMemoryDocumentStore\n", "from haystack.document_stores.types import DuplicatePolicy\n", "\n", "documents = [\n", " Document(\n", " content=\"some publication about Alzheimer prevention research done over 2023 patients study\",\n", " meta={\"year\": 2022, \"disease\": \"Alzheimer\", \"author\": \"Michael Butter\"}),\n", " Document(\n", " content=\"some text about investigation and treatment of Alzheimer disease\",\n", " meta={\"year\": 2023, \"disease\": \"Alzheimer\", \"author\": \"John Bread\"}),\n", " Document(\n", " content=\"A study on the effectiveness of new therapies for Parkinson's disease\",\n", " meta={\"year\": 2022, \"disease\": \"Parkinson\", \"author\": \"Alice Smith\"}\n", " ),\n", " Document(\n", " content=\"An overview of the latest research on the genetics of Parkinson's disease and its implications for treatment\",\n", " meta={\"year\": 2023, \"disease\": \"Parkinson\", \"author\": \"David Jones\"}\n", " )\n", "]\n", "\n", "document_store = InMemoryDocumentStore(bm25_algorithm=\"BM25Plus\")\n", "document_store.write_documents(documents=documents, policy=DuplicatePolicy.OVERWRITE)" ] }, { "cell_type": "markdown", "metadata": { "id": "_7piBFE8AjoG" }, "source": [ "We then create a pipeline consisting of the `QueryMetadataExtractor` and a `InMemoryBM25Retriever` connected to the `InMemoryDocumentStore` created above.\n", "\n", "> Learn about connecting components and creating pipelines in [Docs: Creating Pipelines](https://docs.haystack.deepset.ai/docs/creating-pipelines).\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "J9wS1B92OqhX", "outputId": "cefe4379-c4ac-4ba5-c436-01527005699c" }, "outputs": [ { "data": { "text/plain": [ "\n", "🚅 Components\n", " - metadata_extractor: LLMMetadataQueryExtractor\n", " - retriever: InMemoryBM25Retriever\n", "🛤️ Connections\n", " - metadata_extractor.filters -> retriever.filters (Dict[str, str])" ] }, "execution_count": 18, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from haystack import Pipeline, Document\n", "from haystack.components.retrievers.in_memory import InMemoryBM25Retriever\n", "\n", "\n", "retrieval_pipeline = Pipeline()\n", "metadata_extractor = QueryMetadataExtractor()\n", "retriever = InMemoryBM25Retriever(document_store=document_store)\n", "\n", "retrieval_pipeline.add_component(instance=metadata_extractor, name=\"metadata_extractor\")\n", "retrieval_pipeline.add_component(instance=retriever, name=\"retriever\")\n", "retrieval_pipeline.connect(\"metadata_extractor.filters\", \"retriever.filters\")" ] }, { "cell_type": "markdown", "metadata": { "id": "j-XQ82NrArlS" }, "source": [ "Now define a query and metadata fields and pass them to the pipeline:" ] }, { "cell_type": "code", "execution_count": 24, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 84, "referenced_widgets": [ "e9539a733c51456abfce61c126f8c89e", "129af32510f9453ea6ad319cf15b7d31", "1f4f5294e0b34a88b98ffd6821862d32", "b2d4f5d00e9b4670971eb5ad275ec8cc", "df8553a060d04a1d866c90346ea17189", "b7267bb9646147a49439c2a92eeff4d6", "951f44a228864876987b131643c713e1", "ba46f580af5247ac804e3186a16c32b7", "39c5811e69b44d53ba9ddea9926ff2ac", "40451f91bbc44555ad6d6242be408878", "986c2fa674b94fe29827f7b6bb6e3fc7" ] }, "id": "Jw0XqD7OOsuX", "outputId": "88edf766-9f9d-4627-8bba-7c02bdf64ee2" }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "e9539a733c51456abfce61c126f8c89e", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Ranking by BM25...: 0%| | 0/1 [00:00