feat: Remove Inference from Env Config

Author: kieran-wilkinson-4

.github/scripts/fix_cdk_json.sh

@@ -61,6 +61,3 @@ fix_string_key slackBotToken "${SLACK_BOT_TOKEN}"
 fix_string_key slackSigningSecret "${SLACK_SIGNING_SECRET}"
 fix_string_key cfnDriftDetectionGroup "${CFN_DRIFT_DETECTION_GROUP}"
 fix_boolean_number_key isPullRequest "${IS_PULL_REQUEST}"
-fix_boolean_number_key ragTemperature "${RAG_TEMPERATURE}"
-fix_boolean_number_key ragMaxTokens "${RAG_MAX_TOKENS}"
-fix_boolean_number_key ragTopP "${RAG_TOP_P}"

packages/slackBotFunction/app/core/config.py

@@ -80,9 +80,6 @@ def get_retrieve_generate_config() -> Tuple[str, str, str, str, str, str, str, s
     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"]
-    RAG_TEMPERATURE = os.environ["RAG_TEMPERATURE"]
-    RAG_MAX_TOKENS = os.environ["RAG_MAX_TOKENS"]
-    RAG_TOP_P = os.environ["RAG_TOP_P"]
 
     logger.info(
         "Guardrail configuration loaded", extra={"guardrail_id": GUARD_RAIL_ID, "guardrail_version": GUARD_VERSION}
@@ -95,9 +92,6 @@ def get_retrieve_generate_config() -> Tuple[str, str, str, str, str, str, str, s
         GUARD_VERSION,
         RAG_RESPONSE_PROMPT_NAME,
         RAG_RESPONSE_PROMPT_VERSION,
-        RAG_TEMPERATURE,
-        RAG_MAX_TOKENS,
-        RAG_TOP_P,
     )
 
 

packages/slackBotFunction/app/services/bedrock.py

@@ -1,5 +1,4 @@
 import json
-import os
 from typing import Any
 import boto3
 from mypy_boto3_bedrock_agent_runtime import AgentsforBedrockRuntimeClient
@@ -29,12 +28,18 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
         GUARD_VERSION,
         RAG_RESPONSE_PROMPT_NAME,
         RAG_RESPONSE_PROMPT_VERSION,
-        RAG_TEMPERATURE,
-        RAG_MAX_TOKENS,
-        RAG_TOP_P,
     ) = get_retrieve_generate_config()
 
     prompt_template = load_prompt(RAG_RESPONSE_PROMPT_NAME, RAG_RESPONSE_PROMPT_VERSION)
+    inference_config = prompt_template.get("inference_config")
+
+    if not inference_config:
+        default_values = {"temperature": 0, "maxTokens": 512, "topP": 1}
+        inference_config = default_values
+        logger.warning(
+            "No inference configuration found in prompt template; using default values",
+            extra={"prompt_name": RAG_RESPONSE_PROMPT_NAME, "default_inference_config": default_values},
+        )
 
     client: AgentsforBedrockRuntimeClient = boto3.client(
         service_name="bedrock-agent-runtime",
@@ -54,9 +59,9 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
                     },
                     "inferenceConfig": {
                         "textInferenceConfig": {
-                            "temperature": RAG_TEMPERATURE,
-                            "topP": RAG_TOP_P,
-                            "maxTokens": RAG_MAX_TOKENS,
+                            "temperature": inference_config.get("temperature", 1),
+                            "topP": inference_config.get("topP", 1),
+                            "maxTokens": inference_config.get("maxTokens", 512),
                             "stopSequences": [
                                 "Human:",
                             ],
@@ -70,7 +75,7 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
     if prompt_template:
         request_params["retrieveAndGenerateConfiguration"]["knowledgeBaseConfiguration"]["generationConfiguration"][
             "promptTemplate"
-        ] = {"textPromptTemplate": prompt_template}
+        ] = {"textPromptTemplate": prompt_template.get("prompt_text")}
         logger.info(
             "Using prompt template for RAG response generation", extra={"prompt_name": RAG_RESPONSE_PROMPT_NAME}
         )
@@ -90,16 +95,16 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
     return response
 
 
-def invoke_model(prompt: str, model_id: str, client: BedrockRuntimeClient) -> dict[str, Any]:
+def invoke_model(prompt: str, model_id: str, client: BedrockRuntimeClient, inference_config: dict) -> dict[str, Any]:
     response = client.invoke_model(
         modelId=model_id,
         body=json.dumps(
             {
                 "anthropic_version": "bedrock-2023-05-31",
-                "temperature": os.environ.get("RAG_TEMPERATURE", "1"),
-                "top_p": os.environ.get("RAG_TOP_P", "1"),
+                "temperature": inference_config.get("temperature", "1"),
+                "top_p": inference_config.get("topP", "1"),
                 "top_k": 50,
-                "max_tokens": os.environ.get("RAG_MAX_TOKENS", "512"),
+                "max_tokens": inference_config.get("maxTokens", "512"),
                 "messages": [{"role": "user", "content": prompt}],
             }
         ),

packages/slackBotFunction/app/services/prompt_loader.py

@@ -71,7 +71,7 @@ def parse_system_message(chat_cfg: dict) -> str:
     return "\n\n".join(parts)
 
 
-def load_prompt(prompt_name: str, prompt_version: str = None) -> str:
+def load_prompt(prompt_name: str, prompt_version: str = None) -> dict:
     """
     Load a prompt template from Amazon Bedrock Prompt Management.
 
@@ -103,16 +103,18 @@ def load_prompt(prompt_name: str, prompt_version: str = None) -> str:
         template_config = response["variants"][0]["templateConfiguration"]
         prompt_text = _render_prompt(template_config)
         actual_version = response.get("version", "DRAFT")
+        inference_config = response["variants"][0]["inferenceConfiguration"]
 
         logger.info(
             f"Successfully loaded prompt '{prompt_name}' version {actual_version}",
             extra={
                 "prompt_name": prompt_name,
                 "prompt_id": prompt_id,
                 "version_used": actual_version,
+                "inference_config": inference_config,
             },
         )
-        return prompt_text
+        return {"prompt_text": prompt_text, "inference_config": inference_config}
 
     except ClientError as e:
         error_code = e.response.get("Error", {}).get("Code", "Unknown")

packages/slackBotFunction/app/services/query_reformulator.py

@@ -38,8 +38,10 @@ def reformulate_query(user_query: str) -> str:
         )
 
         # Format the prompt with the user query (using double braces from Bedrock template)
-        prompt = prompt_template.replace("{{user_query}}", user_query)
-        result = invoke_model(prompt=prompt, model_id=model_id, client=client)
+        prompt = prompt_template.get("prompt_text").replace("{{user_query}}", user_query)
+        result = invoke_model(
+            prompt=prompt, model_id=model_id, client=client, inference_config=prompt_template.get("inference_config")
+        )
 
         reformulated_query = result["content"][0]["text"].strip()
 

packages/slackBotFunction/tests/conftest.py

@@ -26,9 +26,6 @@ def mock_env():
         "QUERY_REFORMULATION_PROMPT_VERSION": "DRAFT",
         "RAG_RESPONSE_PROMPT_NAME": "test-rag-prompt",
         "RAG_RESPONSE_PROMPT_VERSION": "DRAFT",
-        "RAG_TEMPERATURE": "0.5",
-        "RAG_MAX_TOKENS": "1024",
-        "RAG_TOP_P": "0.9",
     }
     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

@@ -84,7 +84,7 @@ def test_query_bedrock_check_prompt(mock_boto_client: Mock, mock_load_prompt: Mo
     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"
+    mock_load_prompt.return_value = {"prompt_text": "Test prompt template", "inference_config": {}}
 
     # delete and import module to test
     if "app.services.bedrock" in sys.modules:
@@ -112,6 +112,10 @@ def test_query_bedrock_check_config(mock_boto_client: Mock, mock_load_prompt: Mo
     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 = {
+        "prompt_text": "Test prompt template",
+        "inference_config": {"temperature": "0", "maxTokens": "512", "topP": "1"},
+    }
 
     # delete and import module to test
     if "app.services.bedrock" in sys.modules:
@@ -127,6 +131,6 @@ def test_query_bedrock_check_config(mock_boto_client: Mock, mock_load_prompt: Mo
         "generationConfiguration"
     ]["inferenceConfig"]["textInferenceConfig"]
 
-    assert prompt_config["temperature"] == "0.5"
-    assert prompt_config["maxTokens"] == "1024"
-    assert prompt_config["topP"] == "0.9"
+    assert prompt_config["temperature"] == "0"
+    assert prompt_config["maxTokens"] == "512"
+    assert prompt_config["topP"] == "1"

packages/slackBotFunction/tests/test_prompt_loader.py

@@ -20,7 +20,7 @@ def test_load_prompt_success_draft(mock_boto_client: Mock, mock_env: Mock):
 
     # Mock get_prompt for DRAFT version
     mock_client.get_prompt.return_value = {
-        "variants": [{"templateConfiguration": {"text": {"text": "Test prompt"}}}],
+        "variants": [{"templateConfiguration": {"text": {"text": "Test prompt"}}, "inferenceConfiguration": {}}],
         "version": "DRAFT",
     }
 
@@ -33,7 +33,7 @@ def test_load_prompt_success_draft(mock_boto_client: Mock, mock_env: Mock):
     result = load_prompt("test-prompt")
 
     # assertions
-    assert result == "Test prompt"
+    assert result.get("prompt_text") == "Test prompt"
     mock_client.get_prompt.assert_called_once_with(promptIdentifier="ABC1234567")
 
 
@@ -46,7 +46,7 @@ def test_load_prompt_success_versioned(mock_boto_client: Mock, mock_env: Mock):
     mock_client.list_prompts.return_value = {"promptSummaries": [{"name": "test-prompt", "id": "ABC1234567"}]}
 
     mock_client.get_prompt.return_value = {
-        "variants": [{"templateConfiguration": {"text": {"text": "Versioned prompt"}}}],
+        "variants": [{"templateConfiguration": {"text": {"text": "Versioned prompt"}}, "inferenceConfiguration": {}}],
         "version": "1",
     }
 
@@ -59,7 +59,7 @@ def test_load_prompt_success_versioned(mock_boto_client: Mock, mock_env: Mock):
     result = load_prompt("test-prompt", "1")
 
     # assertions
-    assert result == "Versioned prompt"
+    assert result.get("prompt_text") == "Versioned prompt"
     mock_client.get_prompt.assert_called_once_with(promptIdentifier="ABC1234567", promptVersion="1")
 
 

packages/slackBotFunction/tests/test_query_reformulator.py

@@ -14,7 +14,7 @@ def mock_logger():
 def test_reformulate_query_returns_string(mock_invoke_model: Mock, mock_load_prompt: Mock, mock_env: Mock):
     """Test that reformulate_query returns a string without crashing"""
     # set up mocks
-    mock_load_prompt.return_value = "Test reformat. {{user_query}}"
+    mock_load_prompt.return_value = {"prompt_text": "Test reformat. {{user_query}}", "inference_config": {}}
     mock_invoke_model.return_value = {"content": [{"text": "foo"}]}
 
     # delete and import module to test
@@ -24,6 +24,7 @@ def test_reformulate_query_returns_string(mock_invoke_model: Mock, mock_load_pro
 
     # perform operation
     result = reformulate_query("How do I use EPS?")
+    result = result
 
     # assertions
     # Function should return a string (either reformulated or fallback to original)
@@ -32,7 +33,7 @@ def test_reformulate_query_returns_string(mock_invoke_model: Mock, mock_load_pro
     assert result == "foo"
     mock_load_prompt.assert_called_once_with("test-prompt", "DRAFT")
     mock_invoke_model.assert_called_once_with(
-        prompt="Test reformat. How do I use EPS?", model_id="test-model", client=ANY
+        prompt="Test reformat. How do I use EPS?", model_id="test-model", client=ANY, inference_config={}
     )