Added suggestions. Updated prompt and response format.

app.py

@@ -5,6 +5,7 @@ import math
 import json
 import time
 import re
+import xml.etree.ElementTree as ET
 from botocore.client import Config
 
 kb_id = os.getenv('KNOWLEDGE_BASE_ID')
@@ -15,7 +16,7 @@ initial_message = os.getenv('INITIAL_MESSAGE')
 
 # "amazon.titan-text-express-v1"
 # "amazon.titan-text-premier-v1:0"
-amazon_model_id = "amazon.titan-text-express-v1"
+amazon_model_id = "amazon.titan-text-premier-v1:0"
 
 bedrock_config = Config(connect_timeout=120, read_timeout=120, retries={'max_attempts': 0})
 bedrock_client = boto3.client(
@@ -58,11 +59,28 @@ def get_contexts(retrievalResults):
     return contexts
 
 def clean_text(text):
-    return re.sub(r"(Here are some related questions you might be interested in:|For more information, follow the links provided:).*", "", text, flags=re.DOTALL)
+    return re.sub(r"(For more information, follow the links provided:).*", "", text, flags=re.DOTALL)
+
+def parse_suggestions(response):
+    start_tag = "<SGs>"
+    end_tag = "</SGs>"
+    
+    start_index = response.find(start_tag)
+    end_index = response.find(end_tag) + len(end_tag)
+    
+    if start_index == -1 or end_index == -1:
+        return []
+
+    suggestions_xml = response[start_index:end_index]
+    
+    root = ET.fromstring(suggestions_xml)
+    suggestions = [suggestion.text for suggestion in root.findall('SG')]
+    
+    return suggestions
 
 def get_answer(query, history, temperature, top_p, max_token_count):
     contexts = ""
-    unique_article_ids = []
+    unique_article_ids = [] 
     max_words = math.floor(max_token_count*0.75)
 
     retrievalResults, retrieve_execution_time = retrieve(query)
@@ -70,14 +88,14 @@ def get_answer(query, history, temperature, top_p, max_token_count):
 
     if highest_score > 0.45:
         contexts = get_contexts(retrievalResults)
-        
+
         for result in retrievalResults:
             article_id = result['metadata'].get('article_id')
             if article_id not in unique_article_ids:
                 unique_article_ids.append(article_id)
             if len(unique_article_ids) == 3:
                 break
-    
+
     PROMPT_TEMPLATE = f""" 
     System: You are an intelligent assistant helping users understand and navigate website functionalities.
     Your goal is to provide clear, accurate, and contextually relevant answers based on the information provided.
@@ -109,10 +127,11 @@ def get_answer(query, history, temperature, top_p, max_token_count):
 
     Additionally, only if there is sufficient remaining token capacity, provide 2 or 3 related questions that the user might want to ask next based on the topic.
     If possible, use the following structure for these suggested questions:
-    Here are some related questions you might be interested in:
-    · First related question
-    · Second related question
-    · Third related question
+    <SGs>
+        <SG>First related question</SG>
+        <SG>Second related question</SG>
+        <SG>Third related question</SG>
+    </SGs>
     
     Ensure these suggested questions are brief, relevant, and encourage further exploration on the topic.
     Assistant:
@@ -128,7 +147,7 @@ def get_answer(query, history, temperature, top_p, max_token_count):
     })
 
     kwargs = {
-        "modelId": "amazon.titan-text-premier-v1:0",
+        "modelId": amazon_model_id,
         "contentType": "application/json",
         "accept": "*/*",
         "body": body
@@ -144,19 +163,28 @@ def get_answer(query, history, temperature, top_p, max_token_count):
     response_body = json.loads(response.get('body').read())
     response_text = response_body['results'][0]['outputText']
     
+    suggestions = parse_suggestions(response_text)
+
     if unique_article_ids:
         article_urls_text = "\n\nFor more information, follow the links provided:\n" + "\n".join(
-            f"· https://knowledge.operativeiq.com/articles/{article_id}" for article_id in unique_article_ids)
-        response_text += article_urls_text
+            f"— https://knowledge.operativeiq.com/articles/{article_id}" for article_id in unique_article_ids)
+
+    response_json = {
+        "response_text": response_text.split('<SGs>')[0].strip(),
+        "suggestions": suggestions,
+        "article_urls": article_urls_text
+    }
+    
+    # response_json_str = json.dumps(response_json, indent=4)
+    # print("Response Json:\n", response_json_str)
 
     prompt_and_time = f"""
     Prompt:
     {PROMPT_TEMPLATE}
-
     Retrieve execution time: {retrieve_execution_time} seconds
     Invoke model execution time: {invoke_model_time} seconds
     """
-    return response_text, prompt_and_time
+    return response_json["response_text"] + response_json["article_urls"], prompt_and_time, response_json["suggestions"] 
 
 def format_chat_history(chat_history):
     prompt = ""
@@ -170,23 +198,32 @@ def respond(message, chat_history, temperature=0.9, top_p=0.6, max_token_count=5
     formatted_history = format_chat_history(chat_history[-4:])
     chat_history.append([message, ""])
 
-    stream, prompt_and_time = get_answer(message, formatted_history, temperature, top_p, max_token_count)
-    
+    stream, prompt_and_time, suggestions = get_answer(message, formatted_history, temperature, top_p, max_token_count)
+
     for idx, text_token in enumerate(stream):
         if idx == 0 and text_token.startswith(" "):
             text_token = text_token[1:]
 
         chat_history[-1][1] += text_token
-        yield "", chat_history, prompt_and_time
+        yield "", chat_history, prompt_and_time, suggestions[0], suggestions[1], suggestions[2]
     
 def clear_chat_history():
     return '', []
 
 def main():
-    with gr.Blocks() as demo:
-        chatbot = gr.Chatbot([[None, initial_message]], height=550)
+    with gr.Blocks(css="""
+        .suggestion-button {
+            font-size: 14.5px;
+        }
+    """) as demo:
+        chatbot = gr.Chatbot([[None, initial_message]], height=600)
+
+        with gr.Row():
+            suggestion1 = gr.Button("How to edit RFID crew badge?", elem_classes="suggestion-button")
+            suggestion2 = gr.Button("What types of RFID tags exist?", elem_classes="suggestion-button")
+            suggestion3 = gr.Button("Is it possible to delete a facility?", elem_classes="suggestion-button")
 
-        msg = gr.Textbox(label="Question", elem_id="input_txt")
+        msg = gr.Textbox(label="Question")
         
         with gr.Accordion(label="Advanced options", open=False):
             temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1, value=0.6, step=0.1)
@@ -194,12 +231,37 @@ def main():
             max_token_count = gr.Slider(label="Max token count", minimum=1, maximum=1024, value=400, step=10)
             prompt_and_time = gr.Textbox(label="Prompt and Time", interactive=False)
         
-        btn = gr.Button("Submit", elem_id="btn_submit")
-        clear = gr.Button("Clear console")
+        btn = gr.Button("Submit")
+        clear = gr.Button("Clear history")
 
-        btn.click(respond, inputs=[msg, chatbot, temperature, top_p, max_token_count], outputs=[msg, chatbot, prompt_and_time])
-        msg.submit(respond, inputs=[msg, chatbot, temperature, top_p, max_token_count], outputs=[msg, chatbot, prompt_and_time])
+        btn.click(
+            respond, 
+            inputs=[msg, chatbot, temperature, top_p, max_token_count], 
+            outputs=[msg, chatbot, prompt_and_time, suggestion1, suggestion2, suggestion3]
+        )
 
+        suggestion1.click(lambda s: s, inputs=suggestion1, outputs=msg).then(
+            respond, 
+            inputs=[msg, chatbot, temperature, top_p, max_token_count], 
+            outputs=[msg, chatbot, prompt_and_time, suggestion1, suggestion2, suggestion3]
+        )
+        suggestion2.click(lambda s: s, inputs=suggestion2, outputs=msg).then(
+            respond, 
+            inputs=[msg, chatbot, temperature, top_p, max_token_count], 
+            outputs=[msg, chatbot, prompt_and_time, suggestion1, suggestion2, suggestion3]
+        )
+        suggestion3.click(lambda s: s, inputs=suggestion3, outputs=msg).then(
+            respond, 
+            inputs=[msg, chatbot, temperature, top_p, max_token_count], 
+            outputs=[msg, chatbot, prompt_and_time, suggestion1, suggestion2, suggestion3]
+        )
+    
+        msg.submit(
+            respond, 
+            inputs=[msg, chatbot, temperature, top_p, max_token_count],
+            outputs=[msg, chatbot, prompt_and_time, suggestion1, suggestion2, suggestion3]
+        )
+    
         clear.click(clear_chat_history, outputs=[msg, chatbot, prompt_and_time])
     
     demo.queue().launch()