feat: Add RAG environment variables and update test cases accordingly

Author: kieran-wilkinson-4

packages/slackBotFunction/app/core/config.py

@@ -71,18 +71,28 @@ def get_bot_token() -> str:
 
 
 @lru_cache
-def get_guardrail_config() -> Tuple[str, str, str, str, str]:
+def get_retrieve_generate_config() -> Tuple[str, str, str, str, str, str, str]:
     # 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 (
+        KNOWLEDGEBASE_ID,
+        RAG_MODEL_ID,
+        AWS_REGION,
+        GUARD_RAIL_ID,
+        GUARD_VERSION,
+        RAG_RESPONSE_PROMPT_NAME,
+        RAG_RESPONSE_PROMPT_VERSION,
+    )
 
 
 @dataclass

packages/slackBotFunction/app/services/bedrock.py

@@ -5,7 +5,8 @@
 from mypy_boto3_bedrock_runtime.client import BedrockRuntimeClient
 from mypy_boto3_bedrock_agent_runtime.type_defs import RetrieveAndGenerateResponseTypeDef
 
-from app.core.config import get_guardrail_config, get_logger
+from app.core.config import get_retrieve_generate_config, get_logger
+from app.services.prompt_loader import load_prompt
 
 
 logger = get_logger()
@@ -19,7 +20,18 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
     a response using the configured LLM model with guardrails for safety.
     """
 
-    KNOWLEDGEBASE_ID, RAG_MODEL_ID, AWS_REGION, GUARD_RAIL_ID, GUARD_VERSION = get_guardrail_config()
+    (
+        KNOWLEDGEBASE_ID,
+        RAG_MODEL_ID,
+        AWS_REGION,
+        GUARD_RAIL_ID,
+        GUARD_VERSION,
+        RAG_RESPONSE_PROMPT_NAME,
+        RAG_RESPONSE_PROMPT_VERSION,
+    ) = get_retrieve_generate_config()
+
+    prompt_template = load_prompt(RAG_RESPONSE_PROMPT_NAME, RAG_RESPONSE_PROMPT_VERSION)
+
     client: AgentsforBedrockRuntimeClient = boto3.client(
         service_name="bedrock-agent-runtime",
         region_name=AWS_REGION,
@@ -41,6 +53,14 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
         },
     }
 
+    if prompt_template:
+        request_params["retrieveAndGenerateConfiguration"]["knowledgeBaseConfiguration"]["generationConfiguration"][
+            "promptTemplate"
+        ] = {"textPromptTemplate": prompt_template}
+        logger.info(
+            "Using prompt template for RAG response generation", extra={"prompt_name": RAG_RESPONSE_PROMPT_NAME}
+        )
+
     # Include session ID for conversation continuity across messages
     if session_id:
         request_params["sessionId"] = session_id

packages/slackBotFunction/tests/conftest.py

@@ -24,6 +24,8 @@ def mock_env():
         "QUERY_REFORMULATION_MODEL_ID": "test-model",
         "QUERY_REFORMULATION_PROMPT_NAME": "test-prompt",
         "QUERY_REFORMULATION_PROMPT_VERSION": "DRAFT",
+        "RAG_RESPONSE_PROMPT_NAME": "test-rag-prompt",
+        "RAG_RESPONSE_PROMPT_VERSION": "DRAFT",
     }
     env_vars["AWS_DEFAULT_REGION"] = env_vars["AWS_REGION"]
     with patch.dict(os.environ, env_vars, clear=False):

packages/slackBotFunction/tests/test_bedrock_integration.py

@@ -2,8 +2,9 @@
 from unittest.mock import Mock, patch
 
 
+@patch("app.services.prompt_loader.load_prompt")
 @patch("boto3.client")
-def test_get_bedrock_knowledgebase_response(mock_boto_client: Mock, mock_env: Mock):
+def test_get_bedrock_knowledgebase_response(mock_boto_client: Mock, mock_load_prompt: Mock, mock_env: Mock):
     """Test Bedrock knowledge base integration"""
     # set up mocks
     mock_client = Mock()
@@ -19,13 +20,15 @@ def test_get_bedrock_knowledgebase_response(mock_boto_client: Mock, mock_env: Mo
     result = query_bedrock("test query")
 
     # assertions
+    mock_load_prompt.assert_called_once_with("test-rag-prompt", "DRAFT")
     mock_boto_client.assert_called_once_with(service_name="bedrock-agent-runtime", region_name="eu-west-2")
     mock_client.retrieve_and_generate.assert_called_once()
     assert result["output"]["text"] == "bedrock response"
 
 
+@patch("app.services.prompt_loader.load_prompt")
 @patch("boto3.client")
-def test_query_bedrock_with_session(mock_boto_client: Mock, mock_env: Mock):
+def test_query_bedrock_with_session(mock_boto_client: Mock, mock_load_prompt: Mock, mock_env: Mock):
     """Test query_bedrock with existing session"""
     # set up mocks
     mock_client = Mock()
@@ -42,13 +45,15 @@ def test_query_bedrock_with_session(mock_boto_client: Mock, mock_env: Mock):
     result = query_bedrock("test query", session_id="existing_session")
 
     # assertions
+    mock_load_prompt.assert_called_once_with("test-rag-prompt", "DRAFT")
     assert result == mock_response
     call_args = mock_client.retrieve_and_generate.call_args[1]
     assert call_args["sessionId"] == "existing_session"
 
 
+@patch("app.services.prompt_loader.load_prompt")
 @patch("boto3.client")
-def test_query_bedrock_without_session(mock_boto_client: Mock, mock_env: Mock):
+def test_query_bedrock_without_session(mock_boto_client: Mock, mock_load_prompt: Mock, mock_env: Mock):
     """Test query_bedrock without session"""
     # set up mocks
     mock_client = Mock()
@@ -65,6 +70,35 @@ def test_query_bedrock_without_session(mock_boto_client: Mock, mock_env: Mock):
     result = query_bedrock("test query")
 
     # assertions
+    mock_load_prompt.assert_called_once_with("test-rag-prompt", "DRAFT")
     assert result == mock_response
     call_args = mock_client.retrieve_and_generate.call_args[1]
     assert "sessionId" not in call_args
+
+
+@patch("app.services.prompt_loader.load_prompt")
+@patch("boto3.client")
+def test_query_bedrock_check_prompt(mock_boto_client: Mock, mock_load_prompt: Mock, mock_env: Mock):
+    """Test query_bedrock prompt loading"""
+    # set up mocks
+    mock_client = Mock()
+    mock_boto_client.return_value = mock_client
+    mock_client.retrieve_and_generate.return_value = {"output": {"text": "response"}}
+    mock_load_prompt.return_value = "Test prompt template"
+
+    # delete and import module to test
+    if "app.services.bedrock" in sys.modules:
+        del sys.modules["app.services.bedrock"]
+    from app.services.bedrock import query_bedrock
+
+    # perform operation
+    result = query_bedrock("test query")
+
+    # assertions
+    mock_load_prompt.assert_called_once_with("test-rag-prompt", "DRAFT")
+    call_args = mock_client.retrieve_and_generate.call_args[1]
+    prompt_template = call_args["retrieveAndGenerateConfiguration"]["knowledgeBaseConfiguration"][
+        "generationConfiguration"
+    ]["promptTemplate"]["textPromptTemplate"]
+    assert prompt_template == "Test prompt template"
+    assert result["output"]["text"] == "response"