feat: Increase max tokens

Author: kieran-wilkinson-4

packages/slackBotFunction/app/services/bedrock.py

@@ -25,7 +25,7 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
     inference_config = prompt_template.get("inference_config")
 
     if not inference_config:
-        default_values = {"temperature": 0, "maxTokens": 512, "topP": 1}
+        default_values = {"temperature": 0, "maxTokens": 1500, "topP": 1}
         inference_config = default_values
         logger.warning(
             "No inference configuration found in prompt template; using default values",

packages/slackBotFunction/app/services/prompt_loader.py

@@ -106,7 +106,7 @@ def load_prompt(prompt_name: str, prompt_version: str = None) -> dict:
         actual_version = response.get("version", "DRAFT")
 
         # Extract inference configuration with defaults
-        default_inference = {"temperature": 0, "topP": 1, "maxTokens": 512}
+        default_inference = {"temperature": 0, "topP": 1, "maxTokens": 1500}
         raw_inference = response["variants"][0].get("inferenceConfiguration", {})
         raw_text_config = raw_inference.get("textInferenceConfiguration", {})
         inference_config = {**default_inference, **raw_text_config}

packages/slackBotFunction/app/slack/slack_events.py

@@ -615,70 +615,65 @@ def store_feedback(
 
 
 def open_citation(channel: str, timestamp: str, message: Any, params: Dict[str, Any], client: WebClient) -> None:
-    """
-    Open citation - update/ replace message to include citation content
-    """
+    """Open citation - update/replace message to include citation content"""
     logger.info("Opening citation", extra={"channel": channel, "timestamp": timestamp})
-    # Get Message
     try:
-        # Get citation details
+        # Citation details
         title: str = params.get("title", "No title available.")
         body: str = params.get("body", "No citation text available.")
         source_number: str = params.get("source_number")
+        link: str = params.get("link", "")
 
-        # Remove any citation block (and divider), if it exists
+        # Remove any existing citation block/divider
         blocks = message.get("blocks", [])
-        blocks = [block for block in blocks if block.get("block_id") not in ["citation_block", "citation_divider"]]
-
-        # Add formatting
-        title = f"*{title.replace("\n", "")}*"
-        body = f"> {body.replace("\n", "\n> ")}"
-        link: str = params.get("link", "")
+        blocks = [b for b in blocks if b.get("block_id") not in ["citation_block", "citation_divider"]]
 
-        # Highlight selected citation by updating the button style
-        logger.info("Searching for Citation", extra={"cit_n": source_number, "cit_title": title})
+        # Format text
+        title = f"*{title.replace('\n', '')}*"
+        body = f"> {body.replace('\n', '\n> ')}"
 
-        active_id = None
         current_id = f"cite_{source_number}"
+        selected = False
+
+        # Reset all button styles, then set the clicked one
         for block in blocks:
             if block.get("type") == "actions":
                 for element in block.get("elements", []):
                     if element.get("type") == "button":
                         action_id = element.get("action_id")
-
-                        if element.get("style"):
-                            active_id = action_id
-
-                        if action_id == current_id and active_id != current_id:
-                            logger.info(
-                                "Citation found - set active", extra={"cit_n": source_number, "cit_title": title}
-                            )
-                            element["style"] = "primary"
+                        if action_id == current_id:
+                            # Toggle: if already styled, unselect; else select
+                            if element.get("style") == "primary":
+                                element.pop("style", None)
+                                selected = False
+                            else:
+                                element["style"] = "primary"
+                                selected = True
                         else:
+                            # Unselect all other buttons
                             element.pop("style", None)
 
-        if active_id:
-            # Add citation content before feedback block
+        # If selected, insert citation block before feedback
+        if selected:
             citation_block = {
                 "type": "section",
                 "text": {
                     "type": "mrkdwn",
                     "text": (
                         f"{title}\n\n{body}\n\n<{link}|View Source>"
                         if link and link != "none"
-                        else f"*{title}*\n\n{body}"
+                        else f"{title}\n\n{body}"
                     ),
                 },
                 "block_id": "citation_block",
             }
-
-            # Find index of feedback block to insert before it
-            feedback_block_index = next(
-                (i for i, block in enumerate(blocks) if block.get("block_id") == "feedback-divider"), len(blocks)
+            feedback_index = next(
+                (i for i, b in enumerate(blocks) if b.get("block_id") == "feedback-divider"),
+                len(blocks),
             )
-            blocks.insert(feedback_block_index, citation_block)
+            blocks.insert(feedback_index, citation_block)
 
-        # Update message with new blocks
+        # Update Slack message
         logger.info("Updated message body", extra={"blocks": blocks})
         client.chat_update(channel=channel, ts=timestamp, blocks=blocks)
 

packages/slackBotFunction/tests/test_bedrock_integration.py

@@ -114,7 +114,7 @@ def test_query_bedrock_check_config(mock_boto_client: Mock, mock_load_prompt: Mo
     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"},
+        "inference_config": {"temperature": "0", "maxTokens": "1500", "topP": "1"},
     }
 
     # delete and import module to test
@@ -132,5 +132,5 @@ def test_query_bedrock_check_config(mock_boto_client: Mock, mock_load_prompt: Mo
     ]["inferenceConfig"]["textInferenceConfig"]
 
     assert prompt_config["temperature"] == "0"
-    assert prompt_config["maxTokens"] == "512"
+    assert prompt_config["maxTokens"] == "1500"
     assert prompt_config["topP"] == "1"