more typing

Author: anthony-nhs

packages/slackBotFunction/app/services/dynamo.py

@@ -1,10 +1,12 @@
+from typing import Any
 from app.core.config import get_logger, get_slack_bot_state_table
 from time import time
+from mypy_boto3_dynamodb.type_defs import GetItemOutputTableTypeDef
 
 logger = get_logger()
 
 
-def get_state_information(key):
+def get_state_information(key: str) -> GetItemOutputTableTypeDef:
     start_time = time()
     table = get_slack_bot_state_table()
     is_success = True
@@ -27,7 +29,7 @@ def get_state_information(key):
     return results
 
 
-def store_state_information(item, condition=None):
+def store_state_information(item: dict[str, Any], condition: str = None):
     start_time = time()
     table = get_slack_bot_state_table()
     is_success = True
@@ -52,7 +54,7 @@ def store_state_information(item, condition=None):
         )
 
 
-def update_state_information(key, update_expression, expression_attribute_values):
+def update_state_information(key: str, update_expression: str, expression_attribute_values: dict[str, Any]):
     start_time = time()
     table = get_slack_bot_state_table()
     is_success = True
@@ -76,7 +78,7 @@ def update_state_information(key, update_expression, expression_attribute_values
         )
 
 
-def delete_state_information(pk, sk, condition):
+def delete_state_information(pk: str, sk: str, condition: str):
     start_time = time()
     table = get_slack_bot_state_table()
     is_success = True

packages/slackBotFunction/app/services/prompt_loader.py

@@ -2,11 +2,14 @@
 import traceback
 import boto3
 from botocore.exceptions import ClientError
+from app.core.config import get_logger
 from app.services.exceptions import PromptNotFoundError, PromptLoadError
 from mypy_boto3_bedrock_agent import AgentsforBedrockClient
 
+logger = get_logger()
 
-def load_prompt(logger, prompt_name: str, prompt_version: str = None) -> str:
+
+def load_prompt(prompt_name: str, prompt_version: str = None) -> str:
     """
     Load a prompt template from Amazon Bedrock Prompt Management.
 
@@ -17,7 +20,7 @@ def load_prompt(logger, prompt_name: str, prompt_version: str = None) -> str:
         client: AgentsforBedrockClient = boto3.client("bedrock-agent", region_name=os.environ["AWS_REGION"])
 
         # Get the prompt ID from the name
-        prompt_id = get_prompt_id_from_name(logger, client, prompt_name)
+        prompt_id = get_prompt_id_from_name(client, prompt_name)
         if not prompt_id:
             raise PromptNotFoundError(f"Could not find prompt ID for name '{prompt_name}'")
 
@@ -73,7 +76,7 @@ def load_prompt(logger, prompt_name: str, prompt_version: str = None) -> str:
         raise PromptLoadError(f"Unexpected error loading prompt '{prompt_name}': {e}")
 
 
-def get_prompt_id_from_name(logger, client: AgentsforBedrockClient, prompt_name: str) -> str | None:
+def get_prompt_id_from_name(client: AgentsforBedrockClient, prompt_name: str) -> str | None:
     """
     Get the 10-character prompt ID from the prompt name using ListPrompts.
     """

packages/slackBotFunction/app/services/query_reformulator.py

@@ -2,13 +2,16 @@
 import traceback
 import boto3
 
+from app.core.config import get_logger
 from app.services.bedrock import invoke_model
 from .prompt_loader import load_prompt
 from .exceptions import ConfigurationError
 from mypy_boto3_bedrock_runtime.client import BedrockRuntimeClient
 
+logger = get_logger()
 
-def reformulate_query(logger, user_query: str) -> str:
+
+def reformulate_query(user_query: str) -> str:
     """
     Reformulate user query using Claude Haiku for better RAG retrieval.
 
@@ -27,7 +30,7 @@ def reformulate_query(logger, user_query: str) -> str:
             raise ConfigurationError("QUERY_REFORMULATION_PROMPT_NAME environment variable not set")
 
         # Load prompt with specified version (DRAFT by default)
-        prompt_template = load_prompt(logger, prompt_name, prompt_version)
+        prompt_template = load_prompt(prompt_name, prompt_version)
 
         logger.info(
             "Prompt loaded successfully from Bedrock",

packages/slackBotFunction/app/slack/slack_events.py

@@ -202,7 +202,7 @@ def process_async_slack_event(slack_event_data: Dict[str, Any]):
             return
 
         # Reformulate query for better RAG retrieval
-        reformulated_query = reformulate_query(logger, user_query)
+        reformulated_query = reformulate_query(user_query)
 
         # Check if we have an existing Bedrock conversation session
         session_data = get_conversation_session_data(conversation_key)

packages/slackBotFunction/tests/test_prompt_loader.py

@@ -10,7 +10,7 @@ def mock_logger():
 
 
 @patch("boto3.client")
-def test_load_prompt_success_draft(mock_boto_client: Mock, mock_logger: Mock, mock_env: Mock):
+def test_load_prompt_success_draft(mock_boto_client: Mock, mock_env: Mock):
     # set up mocks
     mock_client = MagicMock()
     mock_boto_client.return_value = mock_client
@@ -30,15 +30,15 @@ def test_load_prompt_success_draft(mock_boto_client: Mock, mock_logger: Mock, mo
     from app.services.prompt_loader import load_prompt
 
     # perform operation
-    result = load_prompt(mock_logger, "test-prompt")
+    result = load_prompt("test-prompt")
 
     # assertions
     assert result == "Test prompt"
     mock_client.get_prompt.assert_called_once_with(promptIdentifier="ABC1234567")
 
 
 @patch("boto3.client")
-def test_load_prompt_success_versioned(mock_boto_client: Mock, mock_logger: Mock, mock_env: Mock):
+def test_load_prompt_success_versioned(mock_boto_client: Mock, mock_env: Mock):
     # set up mocks
     mock_client = MagicMock()
     mock_boto_client.return_value = mock_client
@@ -56,15 +56,15 @@ def test_load_prompt_success_versioned(mock_boto_client: Mock, mock_logger: Mock
     from app.services.prompt_loader import load_prompt
 
     # perform operation
-    result = load_prompt(mock_logger, "test-prompt", "1")
+    result = load_prompt("test-prompt", "1")
 
     # assertions
     assert result == "Versioned prompt"
     mock_client.get_prompt.assert_called_once_with(promptIdentifier="ABC1234567", promptVersion="1")
 
 
 @patch("boto3.client")
-def test_load_prompt_not_found(mock_boto_client: Mock, mock_logger: Mock, mock_env: Mock):
+def test_load_prompt_not_found(mock_boto_client: Mock, mock_env: Mock):
     # set up mocks
     mock_client = MagicMock()
     mock_boto_client.return_value = mock_client
@@ -78,11 +78,11 @@ def test_load_prompt_not_found(mock_boto_client: Mock, mock_logger: Mock, mock_e
 
     # perform operation
     with pytest.raises(Exception, match="Could not find prompt ID"):
-        load_prompt(mock_logger, "nonexistent-prompt")
+        load_prompt("nonexistent-prompt")
 
 
 @patch("boto3.client")
-def test_load_prompt_client_error(mock_boto_client: Mock, mock_logger: Mock, mock_env: Mock):
+def test_load_prompt_client_error(mock_boto_client: Mock, mock_env: Mock):
     # set up mocks
     mock_client = MagicMock()
     mock_boto_client.return_value = mock_client
@@ -99,10 +99,10 @@ def test_load_prompt_client_error(mock_boto_client: Mock, mock_logger: Mock, moc
 
     # perform operation
     with pytest.raises(Exception, match="ValidationException - Invalid prompt"):
-        load_prompt(mock_logger, "test-prompt")
+        load_prompt("test-prompt")
 
 
-def test_get_prompt_id_from_name_success(mock_logger: Mock, mock_env: Mock):
+def test_get_prompt_id_from_name_success(mock_env: Mock):
     # set up mocks
     mock_client = MagicMock()
     mock_client.list_prompts.return_value = {"promptSummaries": [{"name": "test-prompt", "id": "ABC1234567"}]}
@@ -113,13 +113,13 @@ def test_get_prompt_id_from_name_success(mock_logger: Mock, mock_env: Mock):
     from app.services.prompt_loader import get_prompt_id_from_name
 
     # perform operation
-    result = get_prompt_id_from_name(mock_logger, mock_client, "test-prompt")
+    result = get_prompt_id_from_name(mock_client, "test-prompt")
 
     # assertions
     assert result == "ABC1234567"
 
 
-def test_get_prompt_id_from_name_not_found(mock_logger: Mock, mock_env: Mock):
+def test_get_prompt_id_from_name_not_found(mock_env: Mock):
     # set up mocks
     mock_client = MagicMock()
     mock_client.list_prompts.return_value = {"promptSummaries": []}
@@ -130,7 +130,7 @@ def test_get_prompt_id_from_name_not_found(mock_logger: Mock, mock_env: Mock):
     from app.services.prompt_loader import get_prompt_id_from_name
 
     # perform operation
-    result = get_prompt_id_from_name(mock_logger, mock_client, "nonexistent")
+    result = get_prompt_id_from_name(mock_client, "nonexistent")
 
     # assertions
     assert result is None
@@ -147,7 +147,7 @@ def test_get_prompt_id_client_error(mock_logger: Mock, mock_env: Mock):
         del sys.modules["app.services.prompt_loader"]
     from app.services.prompt_loader import get_prompt_id_from_name
 
-    result = get_prompt_id_from_name(mock_logger, mock_client, "test-prompt")
+    result = get_prompt_id_from_name(mock_client, "test-prompt")
 
     # assertions
     assert result is None

packages/slackBotFunction/tests/test_query_reformulator.py

@@ -11,9 +11,7 @@ def mock_logger():
 
 @patch("app.services.prompt_loader.load_prompt")
 @patch("app.services.bedrock.invoke_model")
-def test_reformulate_query_returns_string(
-    mock_invoke_model: Mock, mock_load_prompt: Mock, mock_logger: Mock, mock_env: Mock
-):
+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}}"
@@ -25,7 +23,7 @@ def test_reformulate_query_returns_string(
     from app.services.query_reformulator import reformulate_query
 
     # perform operation
-    result = reformulate_query(mock_logger, "How do I use EPS?")
+    result = reformulate_query("How do I use EPS?")
 
     # assertions
     # Function should return a string (either reformulated or fallback to original)
@@ -39,7 +37,7 @@ def test_reformulate_query_returns_string(
 
 
 @patch("app.services.prompt_loader.load_prompt")
-def test_reformulate_query_prompt_load_error(mock_load_prompt: Mock, mock_logger: Mock, mock_env: Mock):
+def test_reformulate_query_prompt_load_error(mock_load_prompt: Mock, mock_env: Mock):
     # set up mocks
     mock_load_prompt.side_effect = Exception("Prompt not found")
 
@@ -50,17 +48,15 @@ def test_reformulate_query_prompt_load_error(mock_load_prompt: Mock, mock_logger
 
     # perform operation
     original_query = "How do I use EPS?"
-    result = reformulate_query(mock_logger, original_query)
+    result = reformulate_query(original_query)
 
     # assertions
     assert result == original_query
 
 
 @patch("app.services.prompt_loader.load_prompt")
 @patch("app.services.bedrock.invoke_model")
-def test_reformulate_query_bedrock_error(
-    mock_invoke_model: Mock, mock_load_prompt: Mock, mock_logger: Mock, mock_env: Mock
-):
+def test_reformulate_query_bedrock_error(mock_invoke_model: Mock, mock_load_prompt: Mock, mock_env: Mock):
     """Test query reformulation with Bedrock API error"""
     # set up mocks
     mock_load_prompt.return_value = "Reformulate this query: {{user_query}}"
@@ -72,7 +68,7 @@ def test_reformulate_query_bedrock_error(
     from app.services.query_reformulator import reformulate_query
 
     # perform operation
-    result = reformulate_query(mock_logger, "original query")
+    result = reformulate_query("original query")
 
     # assertions
     assert result == "original query"