Update: [AEA-5919] - Split reformulation and rag prompts, versioning and refinement (#169)

## Summary

- Routine Change


### Details

- Update CDK for Prompt Management
- Reformulate user prompt for improved user queries (default value for
now)
   - Set system and user prompt for RAG
   - Set prompt inference for consistent results 
   - Move prompts to .txt files for easier readability
- Update slack bot to read prompts and inference from Bedrock
- Add tests for prompt retrieval

---------

Co-authored-by: Bence Gadanyi <[email protected]>

Author: kieran-wilkinson-4

packages/cdk/prompts/reformulationPrompt.txt

@@ -0,0 +1,5 @@
+Return the user query exactly as provided without any modifications, changes, or reformulations.
+Do not alter, rephrase, or modify the input in any way.
+Simply return: {{user_query}}
+
+User Query: {{user_query}}

packages/cdk/prompts/systemPrompt.txt

@@ -0,0 +1,47 @@
+<SystemInstructions>
+  You are an AI assistant designed to provide helpful information and guidance related to healthcare systems,
+  data integration and user setup.
+  
+  <Requirements>
+    1. Break down the question(s) based on the context
+    2. Examine the information provided in the question(s) or requirement(s).
+    3. Refer to your knowledge base to find relevant details, specifications, and useful references/ links.
+    4. The knowledge base is your source of truth before anything else
+    5. Acknowledge explicit and implicit evidence
+       5a. If no explicit evidence is available, state implicit evidence with a caveat
+    6. Provide critical thinking before replying to make the direction actionable and authoritative
+    7. Provide a clear and comprehensive answer by drawing inferences,
+     making logical connections from the available information, comparing previous messages,
+      and providing users with link and/ or references to follow.
+    8. Be clear in answers, direct actions are preferred (eg., "Check Postcode" &gt; "Refer to documentation") 
+  </Requirements>
+  
+  <Constraints>
+    1. Quotes should be italic
+    2. Document titles and document section names should be bold
+    3. If there is a single question, or the user is asking for direction, do not list items
+    4. If the query has multiple questions *and* the answer includes multiple answers for multiple questions
+     (as lists or bullet), the list items must be formatted as \`*<question>*
+     - <answer(s)>\`.
+      4a. If there are multiple questions in the query, shorten the question to less than 50 characters
+  </Constraints>
+  
+  <Output>
+    - Use Markdown, avoid XML
+    - Structured, informative, and tailored to the specific context of the question. 
+    - Provide evidence to support results
+    - Acknowledging any assumptions or limitations in your knowledge or understanding.
+    - Text structure should be in Markdown
+  </Output>
+  
+  <Tone> 
+    Professional, helpful, authoritative.
+  </Tone>
+  
+  <Examples>
+    <Example1>
+      Q: Should alerts be automated?
+      A: *Section 1.14.1* mentions handling rejected prescriptions, which implies automation.
+    </Example1>
+  </Examples>
+</SystemInstructions>

packages/cdk/prompts/userPrompt.txt

@@ -0,0 +1,7 @@
+- Using your knowledge around the National Health Service (NHS), Electronic Prescription Service (EPS) and the Fast Healthcare Interoperability Resources' (FHIR) onboarding, Supplier Conformance Assessment List (SCAL), APIs, developer guides and error resolution; please answer the following question and cite direct quotes and document sections.
+- If my query is asking for instructions (i.e., "How to...", "How do I...") provide step by steps instructions
+- Do not provide general advice or external instructions
+
+<SearchResults>$search_results$</SearchResults>
+
+<UserQuery>{{user_query}}</UserQuery>`

packages/cdk/resources/BedrockPromptResources.ts

@@ -1,34 +1,60 @@
 import {Construct} from "constructs"
-import {BedrockFoundationModel, Prompt, PromptVariant} from "@cdklabs/generative-ai-cdk-constructs/lib/cdk-lib/bedrock"
+import {
+  BedrockFoundationModel,
+  Prompt,
+  PromptVariant
+} from "@cdklabs/generative-ai-cdk-constructs/lib/cdk-lib/bedrock"
+import {BedrockPromptSettings} from "./BedrockPromptSettings"
 
 export interface BedrockPromptResourcesProps {
   readonly stackName: string
+  readonly settings: BedrockPromptSettings
 }
 
 export class BedrockPromptResources extends Construct {
   public readonly queryReformulationPrompt: Prompt
+  public readonly ragResponsePrompt: Prompt
 
   constructor(scope: Construct, id: string, props: BedrockPromptResourcesProps) {
     super(scope, id)
 
-    const claudeModel = BedrockFoundationModel.ANTHROPIC_CLAUDE_HAIKU_V1_0
-    const promptVariant = PromptVariant.text({
+    const claudeHaikuModel = BedrockFoundationModel.ANTHROPIC_CLAUDE_HAIKU_V1_0
+    const claudeSonnetModel = BedrockFoundationModel.ANTHROPIC_CLAUDE_SONNET_V1_0
+
+    const queryReformulationPromptVariant = PromptVariant.text({
       variantName: "default",
-      model: claudeModel,
+      model: claudeHaikuModel,
       promptVariables: ["topic"],
-      promptText: `Return the user query exactly as provided without any modifications, changes, or reformulations.
-Do not alter, rephrase, or modify the input in any way.
-Simply return: {{user_query}}
-
-User Query: {{user_query}}`
+      promptText: props.settings.reformulationPrompt.text
     })
+
     const queryReformulationPrompt = new Prompt(this, "QueryReformulationPrompt", {
       promptName: `${props.stackName}-queryReformulation`,
       description: "Prompt for reformulating user queries to improve RAG retrieval",
-      defaultVariant: promptVariant,
-      variants: [promptVariant]
+      defaultVariant: queryReformulationPromptVariant,
+      variants: [queryReformulationPromptVariant]
+    })
+
+    const ragResponsePromptVariant = PromptVariant.chat({
+      variantName: "default",
+      model: claudeSonnetModel,
+      promptVariables: ["query", "search_results"],
+      system: props.settings.systemPrompt.text,
+      messages: [props.settings.userPrompt]
+    })
+
+    ragResponsePromptVariant.inferenceConfiguration = {
+      text: props.settings.inferenceConfig
+    }
+
+    const ragPrompt = new Prompt(this, "ragResponsePrompt", {
+      promptName: `${props.stackName}-ragResponse`,
+      description: "Prompt for generating RAG responses with knowledge base context and system instructions",
+      defaultVariant: ragResponsePromptVariant,
+      variants: [ragResponsePromptVariant]
     })
 
     this.queryReformulationPrompt = queryReformulationPrompt
+    this.ragResponsePrompt = ragPrompt
   }
 }

packages/cdk/resources/BedrockPromptSettings.ts

@@ -0,0 +1,73 @@
+import * as fs from "fs"
+import {ChatMessage} from "@cdklabs/generative-ai-cdk-constructs/lib/cdk-lib/bedrock"
+import {Construct} from "constructs"
+import {CfnPrompt} from "aws-cdk-lib/aws-bedrock"
+
+export type BedrockPromptSettingsType = "system" | "user" | "reformulation"
+
+/** BedrockPromptSettings is responsible for loading and providing
+ * the system, user, and reformulation prompts along with their
+ * inference configurations.
+ */
+export class BedrockPromptSettings extends Construct {
+  public readonly systemPrompt: ChatMessage
+  public readonly userPrompt: ChatMessage
+  public readonly reformulationPrompt: ChatMessage
+  public readonly inferenceConfig: CfnPrompt.PromptModelInferenceConfigurationProperty
+
+  /**
+   * @param scope The Construct scope
+   * @param id The Construct ID
+   * @param props BedrockPromptSettingsProps containing optional version info for each prompt
+   */
+  constructor(scope: Construct, id: string) {
+    super(scope, id)
+
+    const systemPromptData = this.getTypedPrompt("system")
+    this.systemPrompt = ChatMessage.assistant(systemPromptData.text)
+
+    const userPromptData = this.getTypedPrompt("user")
+    this.userPrompt = ChatMessage.user(userPromptData.text)
+
+    const reformulationPrompt = this.getTypedPrompt("reformulation")
+    this.reformulationPrompt = ChatMessage.user(reformulationPrompt.text)
+
+    this.inferenceConfig = {
+      temperature: 0,
+      topP: 1,
+      maxTokens: 512,
+      stopSequences: [
+        "Human:"
+      ]
+    }
+  }
+
+  /** Get the latest prompt text from files in the specified directory.
+   * If a version is provided, it retrieves that specific version.
+   * Otherwise, it retrieves the latest version based on file naming.
+   *
+   * @param type The type of prompt (system, user, reformulation)
+   * @returns An object containing the prompt text and filename
+   */
+  private getTypedPrompt(type: BedrockPromptSettingsType)
+  : { text: string; filename: string } {
+    // Read all files in the folder)
+    const path = __dirname + "/../prompts"
+    const files = fs
+      .readdirSync(path)
+
+    if (files.length === 0) {
+      throw new Error("No variant files found in the folder.")
+    }
+
+    const file = files.find(file => file.startsWith(`${type}Prompt`))!
+
+    if (!file) {
+      throw new Error(`No prompt file found for type: ${type}`)
+    }
+
+    const text = fs.readFileSync(`${path}/${file}`, "utf-8")
+
+    return {text, filename: file}
+  }
+}

packages/cdk/resources/Functions.ts

@@ -9,7 +9,6 @@ import {TableV2} from "aws-cdk-lib/aws-dynamodb"
 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 QUERY_REFORMULATION_PROMPT_VERSION = "DRAFT"
 const BEDROCK_KB_DATA_SOURCE = "eps-assist-kb-ds"
 const LAMBDA_MEMORY_SIZE = "265"
 
@@ -33,7 +32,10 @@ export interface FunctionsProps {
   readonly slackBotTokenSecret: Secret
   readonly slackBotSigningSecret: Secret
   readonly slackBotStateTable: TableV2
-  readonly promptName: string
+  readonly reformulationPromptName: string
+  readonly ragResponsePromptName: string
+  readonly reformulationPromptVersion: string
+  readonly ragResponsePromptVersion: string
   readonly isPullRequest: boolean
   readonly mainSlackBotLambdaExecutionRoleArn : string
 }
@@ -66,8 +68,10 @@ export class Functions extends Construct {
         "GUARD_RAIL_ID": props.guardrailId,
         "GUARD_RAIL_VERSION": props.guardrailVersion,
         "SLACK_BOT_STATE_TABLE": props.slackBotStateTable.tableName,
-        "QUERY_REFORMULATION_PROMPT_NAME": props.promptName,
-        "QUERY_REFORMULATION_PROMPT_VERSION": QUERY_REFORMULATION_PROMPT_VERSION
+        "QUERY_REFORMULATION_PROMPT_NAME": props.reformulationPromptName,
+        "RAG_RESPONSE_PROMPT_NAME": props.ragResponsePromptName,
+        "QUERY_REFORMULATION_PROMPT_VERSION": props.reformulationPromptVersion,
+        "RAG_RESPONSE_PROMPT_VERSION": props.ragResponsePromptVersion
       }
     })
 

packages/cdk/stacks/EpsAssistMeStack.ts

@@ -19,6 +19,7 @@ import {BedrockPromptResources} from "../resources/BedrockPromptResources"
 import {S3LambdaNotification} from "../constructs/S3LambdaNotification"
 import {VectorIndex} from "../resources/VectorIndex"
 import {ManagedPolicy, PolicyStatement, Role} from "aws-cdk-lib/aws-iam"
+import {BedrockPromptSettings} from "../resources/BedrockPromptSettings"
 
 export interface EpsAssistMeStackProps extends StackProps {
   readonly stackName: string
@@ -66,9 +67,13 @@ export class EpsAssistMeStack extends Stack {
       stackName: props.stackName
     })
 
+    // Create Bedrock Prompt Collection
+    const bedrockPromptCollection = new BedrockPromptSettings(this, "BedrockPromptCollection")
+
     // Create Bedrock Prompt Resources
     const bedrockPromptResources = new BedrockPromptResources(this, "BedrockPromptResources", {
-      stackName: props.stackName
+      stackName: props.stackName,
+      settings: bedrockPromptCollection
     })
 
     // Create Storage construct first as it has no dependencies
@@ -95,7 +100,8 @@ export class EpsAssistMeStack extends Stack {
       stackName: props.stackName,
       collection: openSearchResources.collection
     })
-    // this dependency ensures the OpenSearch access policy is created before the VectorIndex
+
+    // This dependency ensures the OpenSearch access policy is created before the VectorIndex
     // and deleted after the VectorIndex is deleted to prevent deletion or deployment failures
     vectorIndex.node.addDependency(openSearchResources.deploymentPolicy)
 
@@ -147,7 +153,10 @@ export class EpsAssistMeStack extends Stack {
       slackBotTokenSecret: secrets.slackBotTokenSecret,
       slackBotSigningSecret: secrets.slackBotSigningSecret,
       slackBotStateTable: tables.slackBotStateTable.table,
-      promptName: bedrockPromptResources.queryReformulationPrompt.promptName,
+      reformulationPromptName: bedrockPromptResources.queryReformulationPrompt.promptName,
+      ragResponsePromptName: bedrockPromptResources.ragResponsePrompt.promptName,
+      reformulationPromptVersion: bedrockPromptResources.queryReformulationPrompt.promptVersion,
+      ragResponsePromptVersion: bedrockPromptResources.ragResponsePrompt.promptVersion,
       isPullRequest: isPullRequest,
       mainSlackBotLambdaExecutionRoleArn: mainSlackBotLambdaExecutionRoleArn
     })

packages/slackBotFunction/app/core/config.py

@@ -3,6 +3,7 @@
 Sets up all the AWS and Slack connections we need.
 """
 
+from __future__ import annotations
 from dataclasses import dataclass
 from functools import lru_cache
 import os
@@ -71,18 +72,29 @@ def get_bot_token() -> str:
 
 
 @lru_cache
-def get_guardrail_config() -> Tuple[str, str, str, str, str]:
+def get_retrieve_generate_config() -> BedrockConfig:
     # Bedrock configuration from environment
     KNOWLEDGEBASE_ID = os.environ["KNOWLEDGEBASE_ID"]
     RAG_MODEL_ID = os.environ["RAG_MODEL_ID"]
     AWS_REGION = os.environ["AWS_REGION"]
     GUARD_RAIL_ID = os.environ["GUARD_RAIL_ID"]
     GUARD_VERSION = os.environ["GUARD_RAIL_VERSION"]
+    RAG_RESPONSE_PROMPT_NAME = os.environ["RAG_RESPONSE_PROMPT_NAME"]
+    RAG_RESPONSE_PROMPT_VERSION = os.environ["RAG_RESPONSE_PROMPT_VERSION"]
 
     logger.info(
         "Guardrail configuration loaded", extra={"guardrail_id": GUARD_RAIL_ID, "guardrail_version": GUARD_VERSION}
     )
-    return KNOWLEDGEBASE_ID, RAG_MODEL_ID, AWS_REGION, GUARD_RAIL_ID, GUARD_VERSION
+
+    return BedrockConfig(
+        KNOWLEDGEBASE_ID,
+        RAG_MODEL_ID,
+        AWS_REGION,
+        GUARD_RAIL_ID,
+        GUARD_VERSION,
+        RAG_RESPONSE_PROMPT_NAME,
+        RAG_RESPONSE_PROMPT_VERSION,
+    )
 
 
 @dataclass
@@ -127,6 +139,17 @@ class Constants:
 )
 
 
+@dataclass
+class BedrockConfig:
+    KNOWLEDGEBASE_ID: str
+    RAG_MODEL_ID: str
+    AWS_REGION: str
+    GUARD_RAIL_ID: str
+    GUARD_VERSION: str
+    RAG_RESPONSE_PROMPT_NAME: str
+    RAG_RESPONSE_PROMPT_VERSION: str
+
+
 @dataclass
 class BotMessages:
     EMPTY_QUERY: str