Update query reformulator to use direct invoke_model and remove unused PromptManagement construct

Author: kris-szlapa

packages/cdk/resources/Functions.ts

@@ -7,6 +7,8 @@ import {TableV2} from "aws-cdk-lib/aws-dynamodb"
 
 // Claude model for RAG responses
 const RAG_MODEL_ID = "anthropic.claude-3-sonnet-20240229-v1:0"
+// Claude model for query reformulation
+const QUERY_REFORMULATION_MODEL_ID = "anthropic.claude-3-haiku-20240307-v1:0"
 const BEDROCK_KB_DATA_SOURCE = "eps-assist-kb-ds"
 const LAMBDA_MEMORY_SIZE = "265"
 
@@ -60,6 +62,7 @@ export class Functions extends Construct {
       additionalPolicies: [props.slackBotManagedPolicy],
       environmentVariables: {
         "RAG_MODEL_ID": RAG_MODEL_ID,
+        "QUERY_REFORMULATION_MODEL_ID": QUERY_REFORMULATION_MODEL_ID,
         "KNOWLEDGEBASE_ID": props.knowledgeBaseId || "placeholder",
         "BEDROCK_KB_DATA_SOURCE": BEDROCK_KB_DATA_SOURCE,
         "LAMBDA_MEMORY_SIZE": LAMBDA_MEMORY_SIZE,

packages/cdk/resources/IamResources.ts

@@ -152,15 +152,9 @@ export class IamResources extends Construct {
       resources: [`arn:aws:bedrock:${props.region}:${props.account}:guardrail/*`]
     })
 
-    const slackBotPromptPolicy = new PolicyStatement({
-      actions: [
-        "bedrock:InvokeModel",
-        "bedrock:GetPrompt"
-      ],
-      resources: [
-        `arn:aws:bedrock:${props.region}::foundation-model/${QUERY_REFORMULATION_MODEL_ID}`,
-        `arn:aws:bedrock:${props.region}:${props.account}:prompt/*`
-      ]
+    const slackBotQueryReformulationPolicy = new PolicyStatement({
+      actions: ["bedrock:InvokeModel"],
+      resources: [`arn:aws:bedrock:${props.region}::foundation-model/${QUERY_REFORMULATION_MODEL_ID}`]
     })
 
     const slackBotDynamoDbPolicy = new PolicyStatement({
@@ -196,7 +190,7 @@ export class IamResources extends Construct {
         slackBotSSMPolicy,
         slackBotLambdaPolicy,
         slackBotGuardrailPolicy,
-        slackBotPromptPolicy,
+        slackBotQueryReformulationPolicy,
         slackBotDynamoDbPolicy,
         slackBotKmsPolicy
       ]

packages/cdk/resources/PromptManagement.ts

@@ -11,7 +11,6 @@ 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"
@@ -112,11 +111,6 @@ 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,
@@ -133,10 +127,6 @@ 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/cdk/stacks/EpsAssistMeStack.ts

@@ -8,20 +8,35 @@
 
 def reformulate_query(user_query: str) -> str:
     """
-    Reformulate user query using Bedrock Prompt Management for better RAG retrieval.
+    Reformulate user query using Claude Haiku for better RAG retrieval.
     """
     try:
         client = boto3.client("bedrock-runtime", region_name=os.environ["AWS_REGION"])
-        prompt_arn = os.environ["QUERY_REFORMULATION_PROMPT_ARN"]
+        model_id = os.environ["QUERY_REFORMULATION_MODEL_ID"]
+
+        prompt = f"""You are a query reformulation assistant for the NHS EPS API documentation system.
+
+Reformulate user queries to improve retrieval from a knowledge base containing FHIR NHS EPS API documentation.
+
+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: {user_query}
+
+Reformulated Query:"""
 
         response = client.invoke_model(
-            modelId="anthropic.claude-3-haiku-20240307-v1:0",
+            modelId=model_id,
             body=json.dumps(
                 {
                     "anthropic_version": "bedrock-2023-05-31",
                     "max_tokens": 200,
-                    "prompt": prompt_arn,
-                    "variables": {"query": user_query},
+                    "messages": [{"role": "user", "content": prompt}],
                 }
             ),
         )
@@ -30,18 +45,11 @@ def reformulate_query(user_query: str) -> str:
         reformulated_query = result["content"][0]["text"].strip()
 
         logger.info(
-            "Query reformulated",
-            extra={
-                "original_query": user_query,
-                "reformulated_query": reformulated_query,
-            },
+            "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}",
-            extra={"original_query": user_query, "prompt_arn": os.environ.get("QUERY_REFORMULATION_PROMPT_ARN")},
-        )
+        logger.error(f"Error reformulating query: {e}", extra={"original_query": user_query})
         return user_query  # Fallback to original query