feat: Add additional logs

Author: kieran-wilkinson-4

packages/slackBotFunction/app/services/ai_processor.py

@@ -23,7 +23,7 @@ def process_ai_query(user_query: str, session_id: str | None = None) -> AIProces
 
     logger.info(
         "response from bedrock",
-        extra={"has_citations": len(kb_response.get("citations", [])) > 0, "response_text": kb_response},
+        extra={"response_text": kb_response},
     )
 
     return {

packages/slackBotFunction/app/services/bedrock.py

@@ -90,7 +90,7 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
     response = client.retrieve_and_generate(**request_params)
     logger.info(
         "Got Bedrock response",
-        extra={"session_id": response.get("sessionId"), "has_citations": len(response.get("citations", [])) > 0},
+        extra={"session_id": response.get("sessionId")},
     )
     return response
 

packages/slackBotFunction/app/slack/slack_events.py

@@ -160,12 +160,13 @@ def _create_feedback_blocks(
     # Main response block
     blocks = []
 
-    # Citation buttons
+    # Create citation buttons
     if citations is None or len(citations) == 0:
         logger.info("No citations")
     else:
         action_buttons = []
         for i, citation in enumerate(citations):
+            logger.info("Creating citation", extra={"Citation": citation})
             # Create citation blocks
             # keys = ["source number", "title", "filename", "reference text", "link"]
             title = (
@@ -287,11 +288,9 @@ def process_async_slack_action(body: Dict[str, Any], client: WebClient) -> None:
         timestamp = body["message"]["ts"]
 
         # Check if the action is for a citation
-        if (action_id or "").startswith("cite_"):
-            params = json.loads(action_data)
-
+        if action_id == "cite":
             # Update message to include citation content
-            open_citation(channel_id, timestamp, message, params, client)
+            open_citation(channel_id, timestamp, message, action_data, client)
             return
 
         if message_ts and not is_latest_message(conversation_key=conversation_key, message_ts=message_ts):
@@ -404,15 +403,22 @@ def process_slack_message(event: Dict[str, Any], event_id: str, client: WebClien
         response_text = ai_response["text"]
 
         # Split out citation block if present
-        keys = ["source number", "title", "filename", "reference text", "link"]
-        split = response_text.split("------")
-        if len(split) == 1:
-            citations = []
-        else:
+        # Citations are not returned in the object without using `$output_format_instructions$` which overrides the
+        # system prompt. Instead, pull out and format the citations in the prompt manually
+        prompt_value_keys = ["source number", "title", "filename", "reference text", "link"]
+        split = response_text.split("------")  # Citations are separated by ------
+        citations = []
+        if len(split) != 1:
             response_text = split[0]
             citation_block = split[1]
-            citations = re.split("<cit>", citation_block)[1:] or []
-            citations = [dict(zip(keys, citation.split("|"))) for citation in citations]
+
+            citations = re.compile(r"<cit\b[^>]*>(.*?)</cit>", citation_block, re.DOTALL | re.IGNORECASE).findall(
+                citation_block
+            )
+            if len(citations) > 0:
+                logger.info("Found citation(s)", extra={"Citations": citations})
+                citation_dict = [dict(zip(prompt_value_keys, citation.split("|"))) for citation in citations]
+                logger.info("Parsed citation(s)", extra={"CitationsDict": citation_dict})
 
         # Post the answer (plain) to get message_ts
         post_params = {"channel": channel, "text": response_text}

packages/slackBotFunction/tests/test_slack_handlers.py

@@ -201,3 +201,44 @@ def test_feedback_handler_pull_request(
         pull_request_id="123",
     )
     mock_process_async_slack_action.assert_not_called()
+
+
+@patch("app.slack.slack_events.open_citation")
+@patch("app.utils.handler_utils.extract_session_pull_request_id")
+@patch("app.utils.handler_utils.forward_action_to_pull_request_lambda")
+def test_citation(
+    mock_forward_action_to_pull_request_lambda: Mock,
+    mock_extract_session_pull_request_id: Mock,
+    mock_open_citation: Mock,
+    mock_get_parameter: Mock,
+    mock_env: Mock,
+):
+    """Test citations action handler"""
+    # setup mocks
+    mock_extract_session_pull_request_id.return_value = None
+
+    mock_body = {
+        "user": {"id": "U123"},
+        "actions": [
+            {
+                "action_id": "cite",
+                "value": '{"ck": "conv-key", "ch": "C123", "tt": "123", "mt": "456", "title": "title", "body": "body",'
+                + '"link": "citation_link"}',
+            }
+        ],
+        "channel": {"id": "C123"},
+        "message": {"ts": "123", "blocks": []},
+    }
+    mock_client = Mock()
+
+    # delete and import module to test
+    if "app.slack.slack_handlers" in sys.modules:
+        del sys.modules["app.slack.slack_handlers"]
+    from app.slack.slack_handlers import feedback_handler
+
+    # perform operation
+    feedback_handler(body=mock_body, client=mock_client)
+
+    # assertions
+    mock_open_citation.assert_called_once()
+    mock_forward_action_to_pull_request_lambda.assert_not_called()