Create Bedrock Prompt Management resource for query reformulation

Author: kris-szlapa

packages/cdk/resources/IamResources.ts

@@ -150,6 +150,11 @@ export class IamResources extends Construct {
       resources: [`arn:aws:bedrock:${props.region}:${props.account}:guardrail/*`]
     })
 
+    const slackBotPromptPolicy = new PolicyStatement({
+      actions: ["bedrock:InvokeModel"],
+      resources: [`arn:aws:bedrock:${props.region}:${props.account}:prompt/*`]
+    })
+
     const slackBotDynamoDbPolicy = new PolicyStatement({
       actions: [
         "dynamodb:GetItem",
@@ -183,6 +188,7 @@ export class IamResources extends Construct {
         slackBotSSMPolicy,
         slackBotLambdaPolicy,
         slackBotGuardrailPolicy,
+        slackBotPromptPolicy,
         slackBotDynamoDbPolicy,
         slackBotKmsPolicy
       ]

packages/cdk/resources/PromptManagement.ts

@@ -0,0 +1,46 @@
+import {Construct} from "constructs"
+import {CfnPrompt} from "aws-cdk-lib/aws-bedrock"
+
+export interface PromptManagementProps {
+  readonly stackName: string
+}
+
+export class PromptManagement extends Construct {
+  public readonly queryReformulationPrompt: CfnPrompt
+
+  constructor(scope: Construct, id: string, props: PromptManagementProps) {
+    super(scope, id)
+
+    this.queryReformulationPrompt = new CfnPrompt(this, "QueryReformulationPrompt", {
+      name: `${props.stackName}-query-reformulation`,
+      description: "Reformulates user queries for better RAG retrieval",
+      defaultVariant: "default",
+      variants: [{
+        name: "default",
+        templateType: "TEXT",
+        templateConfiguration: {
+          text: {
+            text: `You are a query reformulation assistant for the NHS EPS (Electronic Prescription Service) API ` +
+              `documentation system.
+
+Your task is to reformulate user queries to improve retrieval from a knowledge base containing FHIR NHS EPS API 
+documentation, onboarding guides, and technical specifications.
+
+Guidelines:
+- Expand abbreviations (EPS = Electronic Prescription Service, FHIR = Fast Healthcare Interoperability Resources)
+- Add relevant technical context (API, prescription, dispensing, healthcare)
+- Convert casual language to technical terminology
+- Include synonyms for better matching
+- Keep the core intent intact
+- Focus on NHS, healthcare, prescription, and API-related terms
+
+User Query: {{query}}
+
+Reformulated Query:`
+          }
+        },
+        modelId: "anthropic.claude-3-haiku-20240307-v1:0"
+      }]
+    })
+  }
+}

packages/cdk/stacks/EpsAssistMeStack.ts

@@ -11,6 +11,7 @@ import {Storage} from "../resources/Storage"
 import {Secrets} from "../resources/Secrets"
 import {OpenSearchResources} from "../resources/OpenSearchResources"
 import {VectorKnowledgeBaseResources} from "../resources/VectorKnowledgeBaseResources"
+import {PromptManagement} from "../resources/PromptManagement"
 import {IamResources} from "../resources/IamResources"
 import {VectorIndex} from "../resources/VectorIndex"
 import {DatabaseTables} from "../resources/DatabaseTables"
@@ -111,6 +112,11 @@ export class EpsAssistMeStack extends Stack {
       endpoint
     })
 
+    // Create Prompt Management resources
+    const promptManagement = new PromptManagement(this, "PromptManagement", {
+      stackName: props.stackName
+    })
+
     // Create VectorKnowledgeBase construct after vector index
     const vectorKB = new VectorKnowledgeBaseResources(this, "VectorKB", {
       stackName: props.stackName,
@@ -127,6 +133,10 @@ export class EpsAssistMeStack extends Stack {
     functions.functions.slackBot.function.addEnvironment("GUARD_RAIL_ID", vectorKB.guardrail.attrGuardrailId)
     functions.functions.slackBot.function.addEnvironment("GUARD_RAIL_VERSION", vectorKB.guardrail.attrVersion)
     functions.functions.slackBot.function.addEnvironment("KNOWLEDGEBASE_ID", vectorKB.knowledgeBase.attrKnowledgeBaseId)
+    functions.functions.slackBot.function.addEnvironment(
+      "QUERY_REFORMULATION_PROMPT_ARN",
+      promptManagement.queryReformulationPrompt.attrArn
+    )
 
     // Create Apis and pass the Lambda function
     const apis = new Apis(this, "Apis", {

packages/slackBotFunction/app.py

@@ -10,6 +10,7 @@
 from aws_lambda_powertools.utilities.parameters import get_parameter
 from aws_lambda_powertools.utilities.typing import LambdaContext
 from botocore.exceptions import ClientError
+from query_reformulator import reformulate_query
 
 
 # Initialize Powertools Logger
@@ -145,7 +146,9 @@ def process_async_slack_event(slack_event_data):
             )
             return
 
-        kb_response = get_bedrock_knowledgebase_response(user_query)
+        # Reformulate query for better RAG retrieval
+        reformulated_query = reformulate_query(user_query)
+        kb_response = get_bedrock_knowledgebase_response(reformulated_query)
         response_text = kb_response["output"]["text"]
 
         client.chat_postMessage(channel=channel, text=response_text, thread_ts=thread_ts)

packages/slackBotFunction/query_reformulator.py

@@ -0,0 +1,40 @@
+import os
+import json
+import boto3
+from aws_lambda_powertools import Logger
+
+logger = Logger(service="queryReformulator")
+
+
+def reformulate_query(user_query: str) -> str:
+    """
+    Reformulate user query using Bedrock Prompt Management for better RAG retrieval.
+    """
+    try:
+        client = boto3.client("bedrock-runtime", region_name=os.environ["AWS_REGION"])
+        prompt_arn = os.environ["QUERY_REFORMULATION_PROMPT_ARN"]
+
+        response = client.invoke_model(
+            modelId="anthropic.claude-3-haiku-20240307-v1:0",
+            body=json.dumps(
+                {
+                    "anthropic_version": "bedrock-2023-05-31",
+                    "max_tokens": 200,
+                    "prompt": prompt_arn,
+                    "variables": {"query": user_query},
+                }
+            ),
+        )
+
+        result = json.loads(response["body"].read())
+        reformulated_query = result["content"][0]["text"].strip()
+
+        logger.info(
+            "Query reformulated", extra={"original_query": user_query, "reformulated_query": reformulated_query}
+        )
+
+        return reformulated_query
+
+    except Exception as e:
+        logger.error(f"Error reformulating query: {e}")
+        return user_query  # Fallback to original query

packages/slackBotFunction/tests/test_query_reformulator.py

@@ -0,0 +1,50 @@
+from unittest.mock import patch, MagicMock
+import json
+from query_reformulator import reformulate_query
+
+
+@patch("query_reformulator.boto3.client")
+@patch.dict(
+    "os.environ",
+    {
+        "AWS_REGION": "eu-west-2",
+        "QUERY_REFORMULATION_PROMPT_ARN": "arn:aws:bedrock:eu-west-2:123456789012:prompt/test-prompt",
+    },
+)
+def test_reformulate_query_success(mock_boto_client):
+    # Mock Bedrock response
+    mock_client = MagicMock()
+    mock_boto_client.return_value = mock_client
+
+    mock_response = {"body": MagicMock()}
+    mock_response["body"].read.return_value = json.dumps(
+        {"content": [{"text": "NHS EPS Electronic Prescription Service API FHIR prescription dispensing"}]}
+    ).encode()
+
+    mock_client.invoke_model.return_value = mock_response
+
+    result = reformulate_query("How do I use EPS?")
+
+    assert result == "NHS EPS Electronic Prescription Service API FHIR prescription dispensing"
+    mock_client.invoke_model.assert_called_once()
+
+
+@patch("query_reformulator.boto3.client")
+@patch.dict(
+    "os.environ",
+    {
+        "AWS_REGION": "eu-west-2",
+        "QUERY_REFORMULATION_PROMPT_ARN": "arn:aws:bedrock:eu-west-2:123456789012:prompt/test-prompt",
+    },
+)
+def test_reformulate_query_fallback_on_error(mock_boto_client):
+    # Mock Bedrock client to raise exception
+    mock_client = MagicMock()
+    mock_boto_client.return_value = mock_client
+    mock_client.invoke_model.side_effect = Exception("Bedrock error")
+
+    original_query = "How do I use EPS?"
+    result = reformulate_query(original_query)
+
+    # Should fallback to original query on error
+    assert result == original_query