Use prompt template replace method instead of format

kris-szlapa · kris-szlapa · commit e42bd42daa48 · 2025-08-25T15:25:53.000Z
diff --git a/packages/slackBotFunction/app/services/prompt_loader.py b/packages/slackBotFunction/app/services/prompt_loader.py
@@ -27,18 +27,18 @@ def load_prompt(prompt_name: str, prompt_version: str = None) -> str:
             raise Exception(f"Could not find prompt ID for name '{prompt_name}'")
 
         # Load the prompt with the specified version
-        if prompt_version == "DRAFT":
-            logger.info(
-                f"Loading DRAFT version of prompt '{prompt_name}' (ID: {prompt_id})",
-                extra={"prompt_name": prompt_name, "prompt_id": prompt_id, "prompt_version": "DRAFT"},
-            )
-            response = client.get_prompt(promptIdentifier=prompt_id)
-        else:
+        if prompt_version and prompt_version != "DRAFT":
             logger.info(
                 f"Loading version {prompt_version} of prompt '{prompt_name}' (ID: {prompt_id})",
                 extra={"prompt_name": prompt_name, "prompt_id": prompt_id, "prompt_version": prompt_version},
             )
             response = client.get_prompt(promptIdentifier=prompt_id, promptVersion=str(prompt_version))
+        else:
+            logger.info(
+                f"Loading DRAFT version of prompt '{prompt_name}' (ID: {prompt_id})",
+                extra={"prompt_name": prompt_name, "prompt_id": prompt_id, "prompt_version": "DRAFT"},
+            )
+            response = client.get_prompt(promptIdentifier=prompt_id)
 
         prompt_text = response["variants"][0]["templateConfiguration"]["text"]["text"]
         actual_version = response.get("version", "DRAFT")
@@ -48,9 +48,7 @@ def load_prompt(prompt_name: str, prompt_version: str = None) -> str:
             extra={
                 "prompt_name": prompt_name,
                 "prompt_id": prompt_id,
-                "version_requested": prompt_version,
-                "version_actual": actual_version,
-                "selection_method": "default" if prompt_version is None else "explicit",
+                "version_used": actual_version,
             },
         )
         return prompt_text
diff --git a/packages/slackBotFunction/app/services/query_reformulator.py b/packages/slackBotFunction/app/services/query_reformulator.py
@@ -33,8 +33,8 @@ def reformulate_query(user_query: str) -> str:
             extra={"prompt_name": prompt_name, "version_used": prompt_version},
         )
 
-        # Format the prompt with the user query
-        prompt = prompt_template.format(user_query=user_query)
+        # Format the prompt with the user query (using double braces from Bedrock template)
+        prompt = prompt_template.replace("{{user_query}}", user_query)
 
         response = client.invoke_model(
             modelId=model_id,
diff --git a/packages/slackBotFunction/tests/test_prompt_loader.py b/packages/slackBotFunction/tests/test_prompt_loader.py
@@ -1,21 +1,106 @@
 import pytest
-from unittest.mock import patch
-from app.services.prompt_loader import load_prompt
+from unittest.mock import patch, MagicMock
+from botocore.exceptions import ClientError
+from app.services.prompt_loader import load_prompt, get_prompt_id_from_name
 
 
-def test_load_prompt_function_exists():
-    """Test that the load_prompt function exists and is callable"""
-    assert callable(load_prompt)
+@patch("app.services.prompt_loader.boto3.client")
+@patch.dict("os.environ", {"AWS_REGION": "eu-west-2"})
+def test_load_prompt_success_draft(mock_boto_client):
+    mock_client = MagicMock()
+    mock_boto_client.return_value = mock_client
 
+    # Mock list_prompts to return prompt ID
+    mock_client.list_prompts.return_value = {"promptSummaries": [{"name": "test-prompt", "id": "ABC1234567"}]}
 
-def test_load_prompt_requires_prompt_name():
-    """Test that load_prompt requires a prompt name parameter"""
-    with pytest.raises(TypeError):
-        load_prompt()  # Should fail without prompt_name
+    # Mock get_prompt for DRAFT version
+    mock_client.get_prompt.return_value = {
+        "variants": [{"templateConfiguration": {"text": {"text": "Test prompt"}}}],
+        "version": "DRAFT",
+    }
 
+    result = load_prompt("test-prompt")
+    assert result == "Test prompt"
+    mock_client.get_prompt.assert_called_once_with(promptIdentifier="ABC1234567")
 
-def test_load_prompt_handles_missing_environment():
-    """Test that load_prompt handles missing AWS_REGION environment variable"""
+
+@patch("app.services.prompt_loader.boto3.client")
+@patch.dict("os.environ", {"AWS_REGION": "eu-west-2"})
+def test_load_prompt_success_versioned(mock_boto_client):
+    mock_client = MagicMock()
+    mock_boto_client.return_value = mock_client
+
+    mock_client.list_prompts.return_value = {"promptSummaries": [{"name": "test-prompt", "id": "ABC1234567"}]}
+
+    mock_client.get_prompt.return_value = {
+        "variants": [{"templateConfiguration": {"text": {"text": "Versioned prompt"}}}],
+        "version": "1",
+    }
+
+    result = load_prompt("test-prompt", "1")
+    assert result == "Versioned prompt"
+    mock_client.get_prompt.assert_called_once_with(promptIdentifier="ABC1234567", promptVersion="1")
+
+
+@patch("app.services.prompt_loader.boto3.client")
+@patch.dict("os.environ", {"AWS_REGION": "eu-west-2"})
+def test_load_prompt_not_found(mock_boto_client):
+    mock_client = MagicMock()
+    mock_boto_client.return_value = mock_client
+
+    mock_client.list_prompts.return_value = {"promptSummaries": []}
+
+    with pytest.raises(Exception, match="Could not find prompt ID"):
+        load_prompt("nonexistent-prompt")
+
+
+@patch("app.services.prompt_loader.boto3.client")
+@patch.dict("os.environ", {"AWS_REGION": "eu-west-2"})
+def test_load_prompt_client_error(mock_boto_client):
+    mock_client = MagicMock()
+    mock_boto_client.return_value = mock_client
+
+    mock_client.list_prompts.return_value = {"promptSummaries": [{"name": "test-prompt", "id": "ABC1234567"}]}
+
+    error = ClientError({"Error": {"Code": "ValidationException", "Message": "Invalid prompt"}}, "GetPrompt")
+    mock_client.get_prompt.side_effect = error
+
+    with pytest.raises(Exception, match="ValidationException - Invalid prompt"):
+        load_prompt("test-prompt")
+
+
+@patch("app.services.prompt_loader.boto3.client")
+@patch.dict("os.environ", {"AWS_REGION": "eu-west-2"})
+def test_get_prompt_id_from_name_success(mock_boto_client):
+    mock_client = MagicMock()
+    mock_client.list_prompts.return_value = {"promptSummaries": [{"name": "test-prompt", "id": "ABC1234567"}]}
+
+    result = get_prompt_id_from_name(mock_client, "test-prompt")
+    assert result == "ABC1234567"
+
+
+@patch("app.services.prompt_loader.boto3.client")
+@patch.dict("os.environ", {"AWS_REGION": "eu-west-2"})
+def test_get_prompt_id_from_name_not_found(mock_boto_client):
+    mock_client = MagicMock()
+    mock_client.list_prompts.return_value = {"promptSummaries": []}
+
+    result = get_prompt_id_from_name(mock_client, "nonexistent")
+    assert result is None
+
+
+@patch("app.services.prompt_loader.boto3.client")
+@patch.dict("os.environ", {"AWS_REGION": "eu-west-2"})
+def test_get_prompt_id_client_error(mock_boto_client):
+    mock_client = MagicMock()
+    error = ClientError({"Error": {"Code": "AccessDenied"}}, "ListPrompts")
+    mock_client.list_prompts.side_effect = error
+
+    result = get_prompt_id_from_name(mock_client, "test-prompt")
+    assert result is None
+
+
+def test_load_prompt_missing_environment():
     with patch.dict("os.environ", {}, clear=True):
         with pytest.raises(Exception):
             load_prompt("test-prompt")
diff --git a/packages/slackBotFunction/tests/test_query_reformulator.py b/packages/slackBotFunction/tests/test_query_reformulator.py
@@ -1,12 +1,10 @@
-from unittest.mock import patch, MagicMock
-import json
+from unittest.mock import patch
 from app.services.query_reformulator import reformulate_query
 
 
-def test_reformulate_query_success():
-    with patch("app.services.query_reformulator.load_prompt") as mock_load_prompt, patch(
-        "app.services.query_reformulator.boto3.client"
-    ) as mock_boto_client, patch.dict(
+def test_reformulate_query_returns_string():
+    """Test that reformulate_query returns a string without crashing"""
+    with patch.dict(
         "os.environ",
         {
             "AWS_REGION": "eu-west-2",
@@ -15,58 +13,35 @@ def test_reformulate_query_success():
         },
     ):
 
-        # Mock prompt loading
-        mock_load_prompt.return_value = "Test prompt template with {user_query}"
-
-        # Mock Bedrock client with proper response
-        mock_client = MagicMock()
-        mock_boto_client.return_value = mock_client
-
-        # Create a simple mock that returns the expected JSON
-        mock_client.invoke_model.return_value = {
-            "body": type(
-                "MockBody",
-                (),
-                {
-                    "read": lambda: json.dumps(
-                        {
-                            "content": [
-                                {"text": "NHS EPS Electronic Prescription Service API FHIR prescription dispensing"}
-                            ]
-                        }
-                    ).encode("utf-8")
-                },
-            )()
-        }
-
         result = reformulate_query("How do I use EPS?")
-
-        # Test that function doesn't crash and returns a string
+        # Function should return a string (either reformulated or fallback to original)
         assert isinstance(result, str)
         assert len(result) > 0
 
 
-@patch("app.services.query_reformulator.load_prompt")
-@patch("app.services.query_reformulator.boto3.client")
-@patch.dict(
-    "os.environ",
-    {
-        "AWS_REGION": "eu-west-2",
-        "QUERY_REFORMULATION_MODEL_ID": "anthropic.claude-3-haiku-20240307-v1:0",
-        "QUERY_REFORMULATION_PROMPT_NAME": "query-reformulation",
-    },
-)
-def test_reformulate_query_fallback_on_error(mock_boto_client, mock_load_prompt):
-    # Mock prompt loading
-    mock_load_prompt.return_value = "Test prompt template with {user_query}"
+def test_reformulate_query_prompt_load_error():
+    with patch("app.services.query_reformulator.load_prompt") as mock_load_prompt, patch.dict(
+        "os.environ",
+        {
+            "AWS_REGION": "eu-west-2",
+            "QUERY_REFORMULATION_MODEL_ID": "anthropic.claude-3-haiku-20240307-v1:0",
+            "QUERY_REFORMULATION_PROMPT_NAME": "query-reformulation",
+        },
+    ):
+
+        mock_load_prompt.side_effect = Exception("Prompt not found")
 
-    # 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)
+        assert result == original_query
 
-    original_query = "How do I use EPS?"
-    result = reformulate_query(original_query)
 
-    # Should fallback to original query on error
-    assert result == original_query
+def test_reformulate_query_missing_prompt_name():
+    with patch.dict(
+        "os.environ",
+        {"AWS_REGION": "eu-west-2", "QUERY_REFORMULATION_MODEL_ID": "anthropic.claude-3-haiku-20240307-v1:0"},
+    ):
+
+        original_query = "test query"
+        result = reformulate_query(original_query)
+        assert result == original_query

Original file line number	Diff line number	Diff line change
`@@ -33,8 +33,8 @@ def reformulate_query(user_query: str) -> str:`
`33`	`33`	`extra={"prompt_name": prompt_name, "version_used": prompt_version},`
`34`	`34`	`)`
`35`	`35`
`36`		`- # Format the prompt with the user query`
`37`		`- prompt = prompt_template.format(user_query=user_query)`
	`36`	`+ # Format the prompt with the user query (using double braces from Bedrock template)`
	`37`	`+ prompt = prompt_template.replace("{{user_query}}", user_query)`
`38`	`38`
`39`	`39`	`response = client.invoke_model(`
`40`	`40`	`modelId=model_id,`