text2ES_prompting_guide.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0a29f818-7dac-48dd-9ca5-8c1965424173",
   "metadata": {},
   "source": [
    "# Prompt Engineering Text-to-QueryDSL Using Claude 3 Models "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8cd487a-81ee-4b9c-9987-440e0a0b5143",
   "metadata": {},
   "source": [
    "---\n",
    "## Introduction\n",
    "\n",
    "In this notebook we step through how to leverage Claude 3 models to create a Text-to-queryDSL. Also we step through how to generate synthetic data for faster testing and development for this example using Claude 3 model. OpenSearch (as well as Elasticsearch) provide a search language called query domain-specific language (DSL) that you can use to search your data. Query DSL is a flexible language with a JSON interface. With query DSL, you need to specify a query in the query parameter of the search.\n",
    "\n",
    "The use case for this notebook is a health social media app. The health app includes different posts on the platform for users as well as user profiles for the health application.\n",
    "\n",
    "This solution is intended to help teams looking for Text-to-queryDSL use cases as well as how to develop test data around existing schemas to test functionality. All the code is intended for development purposes only."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9fc5723-7765-47a9-b829-e62581bd9fa7",
   "metadata": {
    "tags": []
   },
   "source": [
    "--- \n",
    "## Anthropic Claude 3 Model Selection\n",
    "There are multiple different models available on Amazon Bedrock from Anthropic but the main models covered in this notebook come from the Anthropic Claude 3 family:\n",
    "\n",
    "### 1. Claude 3.5 Sonnet\n",
    "- **Description:** Anthropic’s most intelligent and advanced model, Claude 3.5 Sonnet, demonstrates exceptional capabilities across a diverse range of tasks and evaluations while also outperforming Claude 3 Opus.\n",
    "- **MaxTokens:** 200K\n",
    "- **Context Window:** 200K\n",
    "- **Languages:** English, Spanish, Japanese, and multiple other languages.\n",
    "- **Supported Use cases:**  RAG or search & retrieval over vast amounts of knowledge, product recommendations, forecasting, targeted marketing, code generation, quality control, parse text from images.\n",
    "\n",
    "### 2. Claude 3 Sonnet\n",
    "- **Description:** Claude 3 Sonnet is engineered to be dependable for scaled AI deployments across a variety of use cases.\n",
    "- **MaxTokens:** 200K\n",
    "- **Context Window:** 200K\n",
    "- **Languages:** English, Spanish, Japanese, and multiple other languages.\n",
    "- **Supported Use cases:** RAG or search & retrieval over vast amounts of knowledge, product recommendations, forecasting, targeted marketing, code generation, quality control, parse text from images.\n",
    "\n",
    "### 3. Claude 3 Haiku\n",
    "- **Description:** Anthropic’s fastest, most compact model for near-instant responsiveness. It answers simple queries and requests with speed.\n",
    "- **MaxTokens:** 200K\n",
    "- **Context Window:** 200K\n",
    "- **Languages:** English, Spanish, Japanese, and multiple other languages.\n",
    "- **Supported Use cases:** Quick and accurate support in live interactions, translations, content moderation, optimize logistics, inventory management, extract knowledge from unstructured data.\n",
    "\n",
    "We chose from the Claude 3 model family for this workload to compare performance, accuracy and cost. Haiku is faster than Sonnet when running inference on Amazon Bedrock; however, Claude 3 Sonnet and Sonnet 3.5 have improved sustained accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ede8f514-e072-4044-a385-254600c4dbd2",
   "metadata": {},
   "source": [
    "--- \n",
    "\n",
    "## The Approach to text-to-QueryDSL for Development Teams\n",
    "\n",
    "When developing solutions from the ground up, it can be difficult to test when there is no example data present as well as prompting Claude models for the intended tasks. \n",
    "\n",
    "Thus, in this notebook we walk through how to develop synthetic schemas, synthetic data based on the schemas, example questions to ask the synthetic data and then running the query against a Opensearch Serverless collection to understand if the generated query is extracting relevant data from the collection. In this example we are using an Opensearch Serverless Collection as opposed to showing Elasticsearch cluster as many customers have migrated to more modern architecture/services. The approach is similar if you have an existing Elasticsearch cluster.\n",
    "\n",
    "### Few-shot text-to-queryDSL prompting\n",
    "For both syntehtic data creation as well as creating the query DSL based on text, the notebook demonstrates using the few-shot prompting approach and using system prompts. Overall when working with Claude models and building prompts for these models, Anthropic states the golden rule of clear prompting is \"sow your prompt to a colleague, ideally someone who has minimal context on the task, and ask them to follow the instructions. If they’re confused, Claude will likely be too.\"\n",
    "\n",
    "We take this same approach and ensure we are sending specific tasks to the LLM and specific examples to our system prompt to direct the model for each task at hand. For this use case, we start with two schemas \"healthpost_schema\" and \"userprofile_schema\" that our example social media health app is based on. Then we leverage these schemas to create the syntehtic data. We begin with these schemas in the proper formatting to create an index in an Opensearch Cluster. We use the synthetic data to create prompting examples to then pass to the LLM for few-shot prompting. Finally, we ingest this data into two indexes for the Opensearch Collection. To generate the Text-to-queryDSL query, we use few-shot prompting to pass in the examples to give context to the LLM of how we want the syntax of the query to look. \n",
    "\n",
    "The entire flow of this notebook also demonstrates prompt chaining as we are passing in one response from the LLM as the input to another request to the LLM."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f651df95-d5ce-4a1a-8d37-3ccdef908abc",
   "metadata": {},
   "source": [
    "---\n",
    "## Prerequisites\n",
    "1. Use kernel either conda_python3, conda_pytorch_p310 or conda_tensorflow2_p310.\n",
    "2. Install the required packages.\n",
    "3. Access to the Claude3 models on Amazon Bedrock and access to the Converse API.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf62573b-3174-47b5-a6a0-b12e938a725b",
   "metadata": {},
   "source": [
    "## Getting Started\n",
    "### Step 0: Install Dependencies and Import Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ffffd53d-591c-494d-bb23-50f396f0aa19",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install -U awscli -qU --force --quiet --no-warn-conflicts\n",
    "!pip install boto3==1.34.127 -qU --force --quiet --no-warn-conflicts\n",
    "!pip install numpy==1.26.4 -qU --force --quiet --no-warn-conflicts\n",
    "!pip install opensearch-py -qU --force --quiet --no-warn-conflicts\n",
    "!pip install requests-aws4auth -qU --force --quiet --no-warn-conflicts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a763a93-f4c9-46c8-9b33-f7321aa22555",
   "metadata": {},
   "source": [
    "Note: When installing libraries using the pip, you may encounter errors or warnings during the installation process. These are generally not critical and can be safely ignored. However, after installing the libraries, it is recommended to restart the kernel or computing environment you are working in. Restarting the kernel ensures that the newly installed libraries are loaded properly and available for use in your code or workflow.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "56106f97-2b41-467a-af1c-f3a63b109113",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /home/ec2-user/.config/sagemaker/config.yaml\n"
     ]
    }
   ],
   "source": [
    "import boto3\n",
    "from botocore.exceptions import ClientError\n",
    "import json\n",
    "import os\n",
    "from opensearchpy import OpenSearch, RequestsHttpConnection\n",
    "from opensearchpy.helpers import bulk\n",
    "import sagemaker\n",
    "import time\n",
    "import random\n",
    "import re\n",
    "from requests_aws4auth import AWS4Auth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "935461c9-16f1-42af-9f93-8663860e0c2c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'arn:aws:iam::709425451936:role/service-role/AmazonSageMaker-ExecutionRole-20240712T100744'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Setup Bedrock Client\n",
    "bedrock_rt= boto3.client(\n",
    "    service_name='bedrock-runtime'\n",
    ")\n",
    "session = boto3.session.Session()\n",
    "region_name = session.region_name\n",
    "# Create a SageMaker session that will be used when creatin the Opensearch Collection\n",
    "sagemaker_role_arn = sagemaker.get_execution_role()\n",
    "sagemaker_role_arn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f197ff74-6900-4e45-9bc3-924ebcfc2049",
   "metadata": {},
   "source": [
    "### Step 1: Read in the userprofile and healthpost schemas and corresponding index files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "90c48109-9b50-4105-86f0-90e5c6024594",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Read in the userprofile and healthpost schema that are preset that are predefined\n",
    "with open('schemas/userprofile_schema.json', 'r') as file:\n",
    "    userprofile_schema = json.load(file)\n",
    "#open the health posts scheams\n",
    "with open('schemas/healthpost_schema.json', 'r') as file:\n",
    "    healthpost_schema= json.load(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9d3201d-cc6f-409f-83c9-b6c9da074532",
   "metadata": {},
   "source": [
    "We also have a schemas formatted for the Opensearch index created later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8beb4fe9-c130-4aac-9c1a-980620304f53",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "with open('schemas/userprofile_schema_index.json', 'r') as file:\n",
    "    userprofile_schema_index = json.load(file)\n",
    "with open('schemas/healthpost_schema_index.json', 'r') as file:\n",
    "    healthpost_schema_index= json.load(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b8bcf31-8c2c-41fa-bf2c-016d9ce84139",
   "metadata": {
    "tags": []
   },
   "source": [
    "We will need to access all the above information as strings later in the notebook so let's convert the json to strings for later use in the prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "869c4e16-bec4-4cbe-89e1-a5289c6d57f2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#get the correct formatting to send to system prompt\n",
    "userprofile_schema_string= json.dumps(userprofile_schema, indent=2)\n",
    "healthpost_schema_string= json.dumps(healthpost_schema, indent=2)\n",
    "userprofile_index_string= json.dumps(userprofile_schema_index, indent=2)\n",
    "healthpost_index_string= json.dumps(healthpost_schema_index, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3034a6b2-b635-4bbd-a37f-2245de23deb2",
   "metadata": {},
   "source": [
    "### Step 2:  Synthetic Data Generation Based on Schemas\n",
    "Now since we have two schemas for healthposts and userprofiles for the health app, let's learn how to generate synthetic data for the schemas using a foundational model. We will use the generate_data() function throughout out the notebook as the single point of calling the LLM of choice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "df32c1b9-9984-40af-b78c-efc0cdfe4063",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_data(bedrock_rt, model_id, system_prompt, query, inference_config):\n",
    "    \"\"\"\n",
    "    Function to call the Bedrock Converse API\n",
    "    \"\"\"\n",
    "    messages = []\n",
    "    messages.append({\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [{\"text\": query}]\n",
    "        }\n",
    "    )\n",
    "    response = bedrock_rt.converse(\n",
    "        modelId=model_id,\n",
    "        system=system_prompt,\n",
    "        messages=messages,\n",
    "        inferenceConfig=inference_config\n",
    "    )\n",
    "    output = response['output']['message']['content'][0]['text']\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2d8a29d1-9ad8-46b3-88b6-44130ecd07e9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Set up variables to switch the model_id depending on which model you'd like to test. We will start with sonnet 3.5 as default.\n",
    "model_id_sonnet3_5 = \"anthropic.claude-3-5-sonnet-20240620-v1:0\"\n",
    "model_id_sonnet = \"anthropic.claude-3-sonnet-20240229-v1:0\"\n",
    "model_id_haiku = \"anthropic.claude-3-haiku-20240307-v1:0\"\n",
    "model_id = model_id_sonnet\n",
    "#we keep it at 0 because we don't want that much creativity\n",
    "inference_config = {\"temperature\": 0}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caf0a61f-dcb7-4ec1-ba00-45957509948b",
   "metadata": {},
   "source": [
    "We will set up two different messages to pass to the Converse API.\n",
    "1. For health posts \n",
    "2. For user profiles. \n",
    "\n",
    "We separate the user prompt but keep the same system prompt for both calls. This is standard best practices as the user prompt can be subjective to different tasks at hand but system prompt will stay the same for both outputs. We also include a prompt to the LLM to leverage <json> html tags. This will help us later on when we need to convert the string output from the LLM into a readable JSON/dictionary for later use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8555f522-56ac-4b30-a880-56c75b627d5a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "message_healthpost= f\"\"\"Please generate 5 Health Post entries based on the following JSON schema. Health Post Schema:{healthpost_schema_string}.\n",
    "        Please provide the generated data in JSON format, do not include any other information in response other than the JSON outputs.\n",
    "        Put each of these additional dictionaries in separate <json> tags.\"\"\"\n",
    "\n",
    "message_userprofile= f\"\"\"Please generate 5 user profile entries based on the following JSON schema. User Profile Schema: {healthpost_schema_string=}\n",
    "        Please provide the generated data in JSON format, do not include any other information in response other than the JSON outputs\n",
    "         Put each of these additional dictionaries in separate <json> tags\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3d8a3b99-8a14-49eb-9c5e-9a7ba537647b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "system_prompt = [{\"text\": \"\"\"You are an AI assistant tasked with generating synthetic data for a health tech social media platform. You will be provided a JSON schema. Your job is to create realistic, diverse, and consistent sample data entries based on the schema.\n",
    "\n",
    "Follow these guidelines:\n",
    "1. Generate data that adheres strictly to the provided schemas.\n",
    "2. Create diverse and realistic entries, considering various demographics, health conditions, and interests.\n",
    "3. Ensure consistency between User Profile and Health Post data (e.g., usernames, user IDs).\n",
    "4. Use realistic values for all fields, including dates, metrics, and engagement statistics.\n",
    "5. Generate geolocation data for major cities around the world.\n",
    "6. Create a mix of verified and non-verified users/posts.\n",
    "7. Vary the sentiment scores and engagement metrics realistically.\n",
    "8. Include a range of health interests, fitness levels, and medical conditions.\n",
    "9. Generate realistic content for post titles and content fields.\n",
    "\n",
    "Remember to maintain data privacy by not using real people's information. All data should be fictional but plausible.\n",
    "\n",
    "\"\"\"\n",
    "}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "da8aa318-22ae-4ab6-a2c0-acb0a5d54873",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Call the generate_data() function and store in a variable for now. You can loop throguh a few examples of this leveraging the same system prompt\n",
    "health_post_data = generate_data(bedrock_rt, model_id_sonnet, system_prompt, message_healthpost, inference_config)\n",
    "user_profile_data = generate_data(bedrock_rt, model_id_sonnet, system_prompt, message_userprofile, inference_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e199de1-b42b-4352-b8dd-873146cdd71d",
   "metadata": {},
   "source": [
    "Claude doesn't have a formal \"JSON Mode\" with constrained sampling. We can still get reliable JSON from Claude and the code below is pulled from Anthropic's cookbook on how to extract the json. Above in our message we included an xml tag for <json> that allows us to denote when there is a separate output from the LLM and where we woudl need to reformat the response to be in json formatting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4a58cdef-623a-4215-905f-4d46d094f950",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#this function extracts the string object between json tags\n",
    "def extract_json_from_xml(input_string):\n",
    "    # Regular expression to find content inside <json> tags\n",
    "    match = re.search(r'<json>(.*?)</json>', input_string, re.DOTALL)\n",
    "    return match.group(1) if match else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8aa43069-dab4-497f-a9a1-a93fdecea835",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#this function converts into a dictionary\n",
    "def extract_between_tags(tag: str, string: str, strip: bool = False) -> list[str]:\n",
    "    ext_list = re.findall(f\"<{tag}>(.+?)</{tag}>\", string, re.DOTALL)\n",
    "    if strip:\n",
    "        ext_list = [e.strip() for e in ext_list]\n",
    "    return ext_list\n",
    "\n",
    "health_posts_dict = [\n",
    "    json.loads(d)\n",
    "    for d in extract_between_tags(\"json\", health_post_data)\n",
    "]\n",
    "\n",
    "user_profile_dict = [\n",
    "    json.loads(d)\n",
    "    for d in extract_between_tags(\"json\", user_profile_data)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fce6690-51cf-4a51-9dc2-23535590c9df",
   "metadata": {},
   "source": [
    "We will store our synthetic data in the 'data' folder, let's create it and then add the data to individual files for later use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "59340a3e-a1f7-4ed6-b678-257defbcf3ac",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if not os.path.exists('data'):\n",
    "    os.makedirs('data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "164b2aca-59e6-415b-aa04-e74e0d3a2206",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if not os.path.exists('data/healthpost_data.json'):  \n",
    "    with open('data/healthpost_data.json', 'w', encoding='utf-8') as f:\n",
    "        json.dump(health_posts_dict, f, ensure_ascii=False, indent=4)\n",
    "if not os.path.exists('data/userprofile_data.json'):  \n",
    "    with open('data/userprofile_data.json', 'w', encoding='utf-8') as f:\n",
    "        json.dump(user_profile_dict, f, ensure_ascii=False, indent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52717b15-252a-4a78-8424-54cf04b33a65",
   "metadata": {},
   "source": [
    "Now we have our raw data contained in \"health_posts_dict\" and \"user_profile_dict\" variables. Let's move onto the next step to create an Opensearch Collection, Opensearch Index and ingest the data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5544d814-916e-4405-96eb-61fe5fecc325",
   "metadata": {},
   "source": [
    "### Step 3: Create Opensearch Collection and Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1e6ae86f-fba9-4500-924a-102b2626994a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#create initial client for opensearch\n",
    "aoss_client = boto3.client('opensearchserverless')\n",
    "suffix = random.randrange(200, 900)\n",
    "identity = boto3.client('sts').get_caller_identity()['Arn']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb284c1d-2df9-4f25-b2b4-b2b888f0170a",
   "metadata": {},
   "source": [
    "Find code to step through creating Opensearch Collections. Code sampled from here: https://github.com/aws-samples/Cohere-on-AWS/blob/main/cohere-cookbooks/Embeddings/Cohere_Embeddings_Search.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9454cfbd-66ae-442a-92cb-79b7b8fcad80",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_policies_in_oss(es_name, aoss_client, role_arn):\n",
    "    \n",
    "    encryption_policy_name = f\"sample-sp-{suffix}\"\n",
    "    network_policy_name = f\"sample-np-{suffix}\"\n",
    "    access_policy_name = f'sample-ap-{suffix}'\n",
    "\n",
    "    try:\n",
    "        encryption_policy = aoss_client.create_security_policy(\n",
    "            name=encryption_policy_name,\n",
    "            policy=json.dumps(\n",
    "                {\n",
    "                    'Rules': [{'Resource': ['collection/' + es_name],\n",
    "                               'ResourceType': 'collection'}],\n",
    "                    'AWSOwnedKey': True\n",
    "                }),\n",
    "            type='encryption'\n",
    "        )\n",
    "    except Exception as ex:\n",
    "        print(ex)\n",
    "    \n",
    "    try:\n",
    "        network_policy = aoss_client.create_security_policy(\n",
    "            name=network_policy_name,\n",
    "            policy=json.dumps(\n",
    "                [\n",
    "                    {'Rules': [{'Resource': ['collection/' + es_name],\n",
    "                                'ResourceType': 'collection'}],\n",
    "                     'AllowFromPublic': True}\n",
    "                ]),\n",
    "            type='network'\n",
    "        )\n",
    "    except Exception as ex:\n",
    "        print(ex)\n",
    "    \n",
    "    try:\n",
    "        \n",
    "        access_policy = aoss_client.create_access_policy(\n",
    "            name=access_policy_name,\n",
    "            policy=json.dumps(\n",
    "                [\n",
    "                    {\n",
    "                        'Rules': [\n",
    "                            {\n",
    "                                'Resource': ['collection/' + es_name],\n",
    "                                'Permission': [\n",
    "                                    'aoss:CreateCollectionItems',\n",
    "                                    'aoss:DeleteCollectionItems',\n",
    "                                    'aoss:UpdateCollectionItems',\n",
    "                                    'aoss:DescribeCollectionItems'],\n",
    "                                'ResourceType': 'collection'\n",
    "                            },\n",
    "                            {\n",
    "                                'Resource': ['index/' + es_name + '/*'],\n",
    "                                'Permission': [\n",
    "                                    'aoss:CreateIndex',\n",
    "                                    'aoss:DeleteIndex',\n",
    "                                    'aoss:UpdateIndex',\n",
    "                                    'aoss:DescribeIndex',\n",
    "                                    'aoss:ReadDocument',\n",
    "                                    'aoss:WriteDocument'],\n",
    "                                'ResourceType': 'index'\n",
    "                            }],\n",
    "                        'Principal': [identity, role_arn],\n",
    "                        'Description': 'Easy data policy'}\n",
    "                ]),\n",
    "            type='data'\n",
    "        )\n",
    "    except Exception as ex:\n",
    "        print(ex)\n",
    "        \n",
    "    return encryption_policy, network_policy, access_policy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68557ef7-d28f-4d90-8a36-3f14f6e438f4",
   "metadata": {},
   "source": [
    "**note**: **Only run the next cell once. If you run it more than once, will error since the policies already exist**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "eae988a6-7427-4656-9196-f8dadc937b37",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create Collection\n",
    "es_name = f'es-collection-{suffix}'\n",
    "\n",
    "encryption_policy, network_policy, access_policy = create_policies_in_oss(es_name=es_name,\n",
    "                       aoss_client=aoss_client,\n",
    "                       role_arn=sagemaker_role_arn)\n",
    "#the type should be SEARCH, can be changed to VECTORSEARCH if we want a vectorDB\n",
    "collection = aoss_client.create_collection(name=es_name,type='SEARCH')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41503c78-f673-4c94-b76c-7866e486ea8c",
   "metadata": {},
   "source": [
    "**reminder**: only run the above cell ONCE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "14e0dc0a-74b5-41c3-91fb-80e0dfec4680",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "l0pj31ls5ah6qj1cqht0.us-east-1.aoss.amazonaws.com\n"
     ]
    }
   ],
   "source": [
    "#extract the host from the Collection ID to be used \n",
    "collection_id = collection['createCollectionDetail']['id']\n",
    "host = collection_id + '.' + region_name + '.aoss.amazonaws.com'\n",
    "print(host)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6bcc198-fc30-4291-a013-1914ec84f35f",
   "metadata": {},
   "source": [
    "The following code will build the Opensearch client. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6fad7f93-58cb-4c3e-be25-c744ba964201",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "service = 'aoss'\n",
    "credentials= boto3.Session().get_credentials()\n",
    "awsauth = AWS4Auth(credentials.access_key, credentials.secret_key,\n",
    "                   region_name, service, session_token=credentials.token)\n",
    "# Build the OpenSearch client\n",
    "oss_client = OpenSearch(\n",
    "    hosts=[{'host': host, 'port': 443}],\n",
    "    http_auth=awsauth,\n",
    "    use_ssl=True,\n",
    "    verify_certs=True,\n",
    "    connection_class=RequestsHttpConnection,\n",
    "    timeout=300\n",
    ")\n",
    "# It can take up to a minute for data access rules to be enforced\n",
    "time.sleep(60)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad03bb2f-28df-4b9c-ac8a-e2034edc7dc0",
   "metadata": {},
   "source": [
    "**oss_client** is the variable denoted the Opensearch Collection. Now, let's create the two indices for our use case based on the data we read in earlier in the notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "62e5d02f-75eb-4c22-94a4-46d290e08151",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#health post\n",
    "healthpost_index = \"healthpost_index\"\n",
    "healthpost_body = {\n",
    "   \"settings\": {\n",
    "        \"index\": {\n",
    "            \"number_of_shards\": 2,\n",
    "            \"number_of_replicas\": 1\n",
    "        }\n",
    "    },\n",
    "   \"mappings\": healthpost_schema_index['mappings']\n",
    "}\n",
    "\n",
    "#user profile\n",
    "userprofile_index = \"userprofile_index\"\n",
    "userprofile_body = {\n",
    "   \"settings\": {\n",
    "        \"index\": {\n",
    "            \"number_of_shards\": 2,\n",
    "            \"number_of_replicas\": 1\n",
    "        }\n",
    "    },\n",
    "   \"mappings\": userprofile_schema_index['mappings']\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "958d6e43-c0f9-4fc4-b93c-463af05479fb",
   "metadata": {},
   "source": [
    "Now, we ingest the data and let's check the response was received succesfully.\n",
    "**note** sometimes it can take the collection anywhere from a few minutes to 10 minutes to create. If you are getting errors, wait a few more minutes or check status of your Opensearch Collection in AWS console. The status needs to be \"active\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9c28ec01-a980-4522-ba00-856d3e5a3947",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "response received for the create index -> {'acknowledged': True, 'shards_acknowledged': True, 'index': 'healthpost_index'}\n",
      "response received for the create index -> {'acknowledged': True, 'shards_acknowledged': True, 'index': 'userprofile_index'}\n"
     ]
    }
   ],
   "source": [
    "# We would get an index already exists exception if the index already exists, and that is okay. Ignore that error if it occurs\n",
    "try:\n",
    "    response_health = oss_client.indices.create(healthpost_index, body=healthpost_body) \n",
    "    response_user = oss_client.indices.create(userprofile_index, body=userprofile_body)\n",
    "    print(f\"response received for the create index -> {response_health}\")\n",
    "    print(f\"response received for the create index -> {response_user}\")\n",
    "\n",
    "except Exception as e:\n",
    "    print(f\"error, exception={e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83d18f57-99ea-476c-96d9-327e5aa99e72",
   "metadata": {},
   "source": [
    "### Step 4: Ingest Synthetic Data into Opensearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "aabbafde-e4ce-47f9-bcda-9aebc92059aa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#read in the data created from earlier\n",
    "with open('data/healthpost_data.json', 'r') as file:\n",
    "    healthposts_json = json.load(file)\n",
    "with open('data/userprofile_data.json', 'r') as file:\n",
    "    userprofile_json = json.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "dd555d96-6915-4acd-be04-71531136a2c9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def ingest_data(posts, user, healthpost_index, userprofile_index, oss_client):\n",
    "    actions = []\n",
    "    for k1, k2 in zip(healthposts_json, userprofile_json):\n",
    "        actions.append(\n",
    "        {\n",
    "            \"_index\": healthpost_index,\n",
    "            \"_source\": k1\n",
    "        })\n",
    "        actions.append(\n",
    "        {\n",
    "            \"_index\": userprofile_index,\n",
    "            \"_source\": k2\n",
    "        })\n",
    "    success, failed = bulk(oss_client, actions)\n",
    "    print(f\"Successfully indexed {success} documents\")\n",
    "    print(f\"Failed to index {len(failed)} documents\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "93ec6973-7e9f-4826-8e29-cd4990277ab9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successfully indexed 10 documents\n",
      "Failed to index 0 documents\n"
     ]
    }
   ],
   "source": [
    "ingest_data(posts= healthposts_json, user= userprofile_json, healthpost_index=healthpost_index, userprofile_index=userprofile_index, oss_client=oss_client)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d71a747-7e6a-4fa6-9387-eb0b77301013",
   "metadata": {},
   "source": [
    "If we get a succesful statement then we are good to go to the next step to create examples."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d4bdf0b-a266-4126-8494-6a57ecc6a474",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Step 5: Create Examples for the Prompt using LLM and Generate the Query\n",
    "Now, we will be setting up a user prompt and system prompt to pass into our generate_text() function to get the examples to user for the few-shot prompting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "00cb158b-9ac3-4e94-ba74-5dd4f1f6c6c1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "message_questions= \"\"\"Generate 2 natural language questions and the corresponding natural language question based on \n",
    "provided data and scheams.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "8dc2d6ad-4df3-4aba-9291-ef53cd03ead5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "system_prompt_query_generation = [\n",
    "    {\n",
    "        \"text\": f\"\"\"You are an expert query dsl generator. Your task is to take an input question and generate a query dsl to answer\n",
    "        the question. Use the schemas and data below to generate the query. Put the query in json tags <json></json>\n",
    "    Schemas:{userprofile_schema_string} {healthpost_schema_string}\n",
    "    Data:{userprofile_json} {healthposts_json}\n",
    "    \n",
    "    Guidelines:\n",
    "    - Ensure the generated query adheres to DSL query syntax\n",
    "    - Do not created new mappings or other items that aren't included in the provided schemas.\n",
    "    - Think through your answer before answering\n",
    "\n",
    "    Output:\n",
    "    - Only output the generated query with the corresponding generated question. The query should only be the raw query nothing else. Do not include any\n",
    "    headers for the query\"\"\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "77d917ed-569a-4b6f-93e0-a994562ad4fd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "system_prompt_questions= [{\n",
    "    \"text\": f\"\"\"You are an expert query dsl generator. Your task is to provide a query dsl and the corresponding natural language question. Use \n",
    "    the schemas and data below to generate the query and question.   \n",
    "    Schemas:{userprofile_schema_string} {healthpost_schema_string}\n",
    "    Data:{userprofile_json} {healthposts_json}\n",
    "    \n",
    "    Guidelines:\n",
    "    - Ensure the generated query adheres to DSL query syntax\n",
    "    - Do not created new mappings or other items that aren't included in the provided schemas.\n",
    "    - Think through your answer before answering\n",
    "\n",
    "    Output:\n",
    "    - Only output the generated query with the corresponding generated question. The query should only be the raw query nothing else. Do not include any\n",
    "    headers for the query\"\"\"\n",
    "}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "37a5b54b-5f24-4816-aa70-d25a4177cd4f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<json>\n",
      "{\n",
      "  \"query\": {\n",
      "    \"bool\": {\n",
      "      \"must\": [\n",
      "        {\n",
      "          \"match\": {\n",
      "            \"post_type\": \"recipe\"\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"range\": {\n",
      "            \"likes_count\": {\n",
      "              \"gte\": 100\n",
      "            }\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"exists\": {\n",
      "            \"field\": \"media_urls\"\n",
      "          }\n",
      "        }\n",
      "      ]\n",
      "    }\n",
      "  }\n",
      "}\n",
      "</json>\n",
      "\n",
      "Question: Find all recipe posts that have at least 100 likes and include media URLs.\n",
      "\n",
      "<json>\n",
      "{\n",
      "  \"query\": {\n",
      "    \"bool\": {\n",
      "      \"must\": [\n",
      "        {\n",
      "          \"match\": {\n",
      "            \"category\": \"fitness\"\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"range\": {\n",
      "            \"health_metrics.steps\": {\n",
      "              \"gte\": 10000\n",
      "            }\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"range\": {\n",
      "            \"health_metrics.heart_rate\": {\n",
      "              \"gte\": 150\n",
      "            }\n",
      "          }\n",
      "        }\n",
      "      ]\n",
      "    }\n",
      "  }\n",
      "}\n",
      "</json>\n",
      "\n",
      "Question: Find all fitness-related posts where the user recorded at least 10,000 steps and had a heart rate of 150 or higher.\n"
     ]
    }
   ],
   "source": [
    "example_prompt= generate_data(bedrock_rt, model_id_sonnet, system_prompt_query_generation, message_questions, inference_config)\n",
    "print(example_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5af2129-f297-4313-9632-bac96c993591",
   "metadata": {},
   "source": [
    "Now, let's look at what the example looks like with the corresponding question and query"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "007813e2-8dae-48ed-a1b0-ad3b0a50f934",
   "metadata": {},
   "source": [
    "Now we are ready to configure the system prompt to help us create the es query. We will provide:\n",
    "1. The schemas\n",
    "2. The index names\n",
    "3. The examples we've generated\n",
    "\n",
    "We add \"Guidelines\" and \"Output\" fields to control the response of the LLM a bit more. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "026cfd70-3499-4ea3-82ba-7b1ae1217cd1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the system prompt for generating queries\n",
    "system_prompt_query_generation = [\n",
    "    {\n",
    "        \"text\": f\"\"\"You are an expert query dsl generator. Your task is to take an input question and generate a query dsl to answer\n",
    "        the question. Use the schemas and data below to generate the query. Put the query in json tags <json></json>\n",
    "    Schemas:{userprofile_schema_string} {healthpost_schema_string}\n",
    "    Data:{userprofile_json} {healthposts_json}\n",
    "    \n",
    "    Examples: {example_prompt}\n",
    "    \n",
    "    Guidelines:\n",
    "    - Ensure the generated query adheres to DSL query syntax\n",
    "    - Do not created new mappings or other items that aren't included in the provided schemas.\n",
    "    - Think through your answer before answering\n",
    "\n",
    "    Output:\n",
    "    - Only output the generated query with the corresponding generated question. The query should only be the raw query nothing else. Do not include any\n",
    "    headers for the query\"\"\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2e1c142-df3c-47c5-bc6d-7422042e7875",
   "metadata": {
    "tags": []
   },
   "source": [
    "Let's pass in an example question to generate the query, then test it against our Opensearch Collection. Remember, we chose Claude 3 Sonnet model as default in the beginning. To test the other models, change out \"model_id\" when calling generate_data()."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "7253f8d6-a74f-47bb-a96c-ffb6f1e4eecf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "query =\"What are some healthy recipes that include quinoa?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "26027479-fd9b-4efa-8bb1-11b8e877e47d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sonnet took 7.273151397705078 to run\n",
      "Response from Sonnet:\n",
      "<json>\n",
      "{\n",
      "  \"query\": {\n",
      "    \"bool\": {\n",
      "      \"must\": [\n",
      "        {\n",
      "          \"match\": {\n",
      "            \"post_type\": \"recipe\"\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"match\": {\n",
      "            \"content\": \"quinoa\"\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"match\": {\n",
      "            \"tags\": \"healthy\"\n",
      "          }\n",
      "        }\n",
      "      ]\n",
      "    }\n",
      "  }\n",
      "}\n",
      "</json>\n"
     ]
    }
   ],
   "source": [
    "# Generate responses and measure runtime with Haiku model\n",
    "start_time = time.time()\n",
    "query_response = generate_data(bedrock_rt, model_id, system_prompt_query_generation, query, inference_config)\n",
    "print(\"Sonnet took\", time.time() - start_time, \"to run\")\n",
    "print(\"Response from Sonnet:\")\n",
    "print(query_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b126a3d7-4112-4573-a844-02e51c581a3c",
   "metadata": {},
   "source": [
    "Define a function for running the generated query against the oss client:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "5106db0f-eba1-4f05-a3e5-9f8a5c9f18cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def query_oss(query):\n",
    "    #extract the json tags with the function generated beforehand\n",
    "    dict1= [\n",
    "    json.loads(d)\n",
    "    for d in extract_between_tags(\"json\", query)\n",
    "    ]\n",
    "    temp = dict1[0]\n",
    "    response = oss_client.search(\n",
    "    index = \"_all\",\n",
    "    body= temp\n",
    "    )\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "c6f8739c-53d4-4cb7-a34e-348a3a5fa21e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'took': 45,\n",
       " 'timed_out': False,\n",
       " '_shards': {'total': 0, 'successful': 0, 'skipped': 0, 'failed': 0},\n",
       " 'hits': {'total': {'value': 1, 'relation': 'eq'},\n",
       "  'max_score': 2.397612,\n",
       "  'hits': [{'_index': 'userprofile_index',\n",
       "    '_id': '7u-9m5EBiT38F8xt6mpC',\n",
       "    '_score': 2.397612,\n",
       "    '_source': {'user_id': 'user_1',\n",
       "     'username': 'FitFoodie123',\n",
       "     'post_type': 'recipe',\n",
       "     'title': 'Healthy Quinoa Salad with Roasted Veggies',\n",
       "     'content': 'Looking for a nutritious and delicious meal? Try this quinoa salad packed with roasted veggies like bell peppers, zucchini, and onions. Drizzle with a tangy lemon vinaigrette for an extra burst of flavor!',\n",
       "     'created_at': '2023-05-01T12:34:56Z',\n",
       "     'updated_at': '2023-05-01T12:34:56Z',\n",
       "     'tags': ['healthy', 'vegetarian', 'meal-prep'],\n",
       "     'category': 'food',\n",
       "     'likes_count': 125,\n",
       "     'comments_count': 32,\n",
       "     'shares_count': 18,\n",
       "     'media_urls': ['https://example.com/quinoa-salad.jpg'],\n",
       "     'location': {'lat': 40.7128, 'lon': -74.0059},\n",
       "     'mentioned_users': [],\n",
       "     'health_metrics': {'steps': 8500,\n",
       "      'calories': 1800.5,\n",
       "      'heart_rate': 72,\n",
       "      'blood_pressure': {'systolic': 120, 'diastolic': 80},\n",
       "      'sleep_duration': 7.2},\n",
       "     'sentiment_score': 0.85,\n",
       "     'is_verified': True}}]}}"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#query the oss client and evaluate results\n",
    "query_oss(query_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2a3460b-0246-497f-86e5-b2f21e3c0aa5",
   "metadata": {},
   "source": [
    "Above we can see that the model was able to generate a query as well as accurately return data from the Opensearch Collection."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25ed6440-dc73-40f8-b015-416509794af1",
   "metadata": {},
   "source": [
    "**note**: if you receive authorization errors eventually, just scroll up and rerun the cell that builds the Opensearch client. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91038799-de98-4d91-9264-6c7747166cb3",
   "metadata": {},
   "source": [
    "### Step 6: Generate Test Questions with LLM\n",
    "In order to improve the amount of tests we have for this synthetic approach let's use the LLM to generate example questions we can start testing. This is helpful when evaluating the performance of the LLM and accuracy of responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "ee3c8fba-d294-453e-847b-f50564d825e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#User prompt for generating the example questions\n",
    "query_prompt= \"\"\"Your task is to generate 5 example questions users can ask the health app based on provided schemas and data. Only include the questions generated\n",
    "        in the response.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "73fd3805-9714-4cd0-8ea3-1b100abe36cc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the system prompt for generating queries\n",
    "query_system_prompt = [\n",
    "    {\n",
    "        \"text\": f\"\"\"Your task is to use the provided schemas and data to generate the questions. The questions you create, should be answered by the data provided. \n",
    "        Do not generate questions that cannot be answered by the data provided.\n",
    "Schemas:\n",
    "{userprofile_schema_string}\n",
    "{healthpost_schema_string}\n",
    "\n",
    "Data:{userprofile_json} {healthposts_json}\n",
    "\n",
    "Guidelines:\n",
    "- Create only the questions that can be answered by the provided information.\n",
    "Output:\n",
    "- Only output the questions, nothing else\n",
    "\n",
    "Role:\n",
    "- Think through your answer\n",
    "\"\"\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "49fe92eb-2589-45cc-8f90-140701dc2b2a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. What are some healthy recipes that include quinoa?\n",
      "2. Which users have shared their running or training experiences?\n",
      "3. Are there any posts discussing mindfulness techniques for busy parents?\n",
      "4. What gluten-free and dairy-free dessert recipes have been shared?\n",
      "5. Are there any yoga or meditation posts that focus on flexibility and strength?\n"
     ]
    }
   ],
   "source": [
    "test_queries = generate_data(bedrock_rt, model_id_sonnet, query_system_prompt, query_prompt, inference_config)\n",
    "print(test_queries)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0cf8b80-a487-47f1-9345-ec34df4fa996",
   "metadata": {},
   "source": [
    "Let's select one of these generated questions. Copy and paste a generate question into the below cell. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "c34048c7-fca0-46d8-9068-bf4d9385bf19",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "user =\"What gluten-free and dairy-free dessert recipes have been shared?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "c39b1636-9da6-47c1-9f9e-c9510a126ba5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<json>\n",
      "{\n",
      "  \"query\": {\n",
      "    \"bool\": {\n",
      "      \"must\": [\n",
      "        {\n",
      "          \"match\": {\n",
      "            \"post_type\": \"recipe\"\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"match\": {\n",
      "            \"category\": \"food\"\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"match\": {\n",
      "            \"tags\": \"gluten-free\"\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"match\": {\n",
      "            \"tags\": \"dairy-free\"\n",
      "          }\n",
      "        }\n",
      "      ]\n",
      "    }\n",
      "  }\n",
      "}\n",
      "</json>\n",
      "Sonnet took 4.676537752151489 to run\n",
      "Response after running query against Opensearch\n",
      "{'took': 47, 'timed_out': False, '_shards': {'total': 0, 'successful': 0, 'skipped': 0, 'failed': 0}, 'hits': {'total': {'value': 1, 'relation': 'eq'}, 'max_score': 5.2397428, 'hits': [{'_index': 'userprofile_index', '_id': '9O-9m5EBiT38F8xt6mpC', '_score': 5.2397428, '_source': {'user_id': 'user_4', 'username': 'GlutenFreeFoodie', 'post_type': 'recipe', 'title': 'Gluten-Free Chocolate Chip Cookies (Dairy-Free Option)', 'content': \"Craving something sweet but need a gluten-free and dairy-free option? These chocolate chip cookies are sure to satisfy your cravings! They're soft, chewy, and absolutely delicious. Get the recipe...\", 'created_at': '2023-05-04T20:38:12Z', 'updated_at': '2023-05-04T20:38:12Z', 'tags': ['gluten-free', 'dairy-free', 'dessert'], 'category': 'food', 'likes_count': 95, 'comments_count': 18, 'shares_count': 25, 'media_urls': ['https://example.com/gf-cookies.jpg'], 'location': {'lat': 19.4326, 'lon': -99.1332}, 'mentioned_users': [], 'health_metrics': {}, 'sentiment_score': 0.91, 'is_verified': False}}]}}\n"
     ]
    }
   ],
   "source": [
    "# Generate responses and measure runtime with Haiku model\n",
    "start_time = time.time()\n",
    "response1 = generate_data(bedrock_rt, model_id_sonnet, system_prompt_query_generation, user, inference_config)\n",
    "print(response1)\n",
    "output = query_oss(response1)\n",
    "print(\"Sonnet took\", time.time() - start_time, \"to run\")\n",
    "print(\"Response after running query against Opensearch\")\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37034e6a-9d3b-42d6-82db-b09543b238a1",
   "metadata": {},
   "source": [
    "Above we can see that the query generated is accurate syntax to query data within Opensearch Collection."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71847015-5d3d-4045-b920-9bb92f9d9558",
   "metadata": {},
   "source": [
    "---\n",
    "## Clean Up\n",
    "\n",
    "After we are done, delete the indexes for the collection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08d68a7e-a942-45e2-b993-5c70988dbdeb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def delete_opensearch_serverless_indices(collection_id, client):\n",
    "    # Create a boto3 client for OpenSearch Serverles\n",
    "    try:\n",
    "        client.indices.delete(index='_all')\n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred: {e}\")\n",
    "        \n",
    "delete_opensearch_serverless_indices(collection_id, oss_client)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09a97084-3bd9-49f7-b195-ee97d6b86411",
   "metadata": {},
   "source": [
    "---\n",
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6a369a-a960-4b54-a506-c3593e15e22e",
   "metadata": {},
   "source": [
    "We observed through this notebook how to create synthetic data to improve pace of testing with schemas and an approach to text-to-queryDSL using Claude 3 models.\n",
    "\n",
    "This notebook allows you to change the model_id to Haiku, Sonnet and Sonnet 3.5 to test accuracy and latency between all three models."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}