Update app.py

app.py

@@ -102,110 +102,111 @@ def ask_ds(message, history):
     global extractions
     
     if len(message) == 0:
-        return
+        yield None
 
     if authenticated == False:
         
         if division == None:
             if message.lower().strip() in ['ime', 'peer disability', 'pas']:
                 division = message.lower().strip().replace(' ', '_')
-                return "[1] CSR\n[2] QA"
+                yield "[1] CSR\n[2] QA"
             else:
-                return "Please select a valid choice."
+                yield "Please select a valid choice."
         elif employee_type == None:
             if message.lower().strip() in ['csr', 'qa']:
                 employee_type = message.lower().strip()
                 authenticated = True
                 EXTRACTIONS_PATH = EXTRACTIONS_PATH.replace('{employee_type}', employee_type).replace('{division}', division[:3])
                 extractions = read_json_from_s3()
-                return "Welcome to Ask Dane Street! Whether you're new to the team or just looking for some quick information, I'm here to guide you through our company's literature and platform. Simply ask your question, and I'll provide you with the most relevant information I can."
+                yield "Welcome to Ask Dane Street! Whether you're new to the team or just looking for some quick information, I'm here to guide you through our company's literature and platform. Simply ask your question, and I'll provide you with the most relevant information I can."
             else:
-                return "Please select a valid choice."    
-
-    question = message
-
-    # RAG
-    question_embedding = get_titan_embedding(bedrock_client, 'question', question)
-    
-    similar_documents = []
-    for file, data in extractions.items():
-        similarity = cosine_similarity(question_embedding, np.array(data['embedding']))
-        similar_documents.append((file, similarity))
-
-    similar_documents.sort(key=lambda x: x[1], reverse=False)
-    top_docs = similar_documents[:5]
-
-    similar_content = ''
-    for file, _ in top_docs:
-        similar_content += extractions[file]['content'] + '\n'
-
-    # Invoke
-    response = bedrock_client.invoke_model_with_response_stream(
-        modelId="anthropic.claude-3-sonnet-20240229-v1:0",
-        body=json.dumps(
-            {
-                "anthropic_version": "bedrock-2023-05-31",
-                "max_tokens": 4096,
-                "system": f"""Here is some relevant information that may help answer the user's upcoming question:
-
-                <relevant_information>
-                {similar_content}
-                </relevant_information>
-                
-                The user's question is:
-                <question>{question}</question>
-
-                Please carefully review the relevant information provided above. 
-
-                Your task is to review the provided relevant information and answer the user's question to the best of your ability. 
-                Aim to use information from the relevant information section to directly address the question asked, and refrain from saying 
-                things like 'According to the relevant information provided'.
-
-                Format your output nicely with sentences that are not too long, in a professional and kind tone. You should prefer lists or 
-                bullet points when applicable. Begin by thanking the user for their question, and at the end of your answer, say "Thank you for using Ask Dane Street!"
-                Remember, aim to only use information from the relevant information section in your response, without explicitly referring 
-                to that section. Return your answer immediately and without preamble. 
-                </Task>
-                
-                <Relevant Information>
-                {similar_content}
-                </Relevant Information>""",
-                "messages": [
-                        {
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "text",
-                                "text": message
-                            }
-                        ]
-                    }
-                ],
-            }
-        ),
-    )
+                yield "Please select a valid choice."   
+    else:
 
-    # Stream the response
-    all_text = ''
-    stream = response.get('body')
-    if stream:
-        for event in stream:
-            chunk = event.get('chunk')
-            if chunk and json.loads(chunk.get('bytes').decode()):
-                # check if delta is present
-                try:
-                    this_text = json.loads(chunk.get('bytes').decode()).get('delta').get('text')
-                    all_text += this_text
-                    yield all_text  # Stream the text back to the UI
-                except:
-                    pass
-
-    # Print relevant files
-    output = '\n\nCheck out the following documents for more information:\n'
-    for file, sim in top_docs:
-        output += f"\n{file.replace('.txt', '.pdf')}"
-
-    yield all_text + output
+        question = message
+    
+        # RAG
+        question_embedding = get_titan_embedding(bedrock_client, 'question', question)
+        
+        similar_documents = []
+        for file, data in extractions.items():
+            similarity = cosine_similarity(question_embedding, np.array(data['embedding']))
+            similar_documents.append((file, similarity))
+    
+        similar_documents.sort(key=lambda x: x[1], reverse=False)
+        top_docs = similar_documents[:5]
+    
+        similar_content = ''
+        for file, _ in top_docs:
+            similar_content += extractions[file]['content'] + '\n'
+    
+        # Invoke
+        response = bedrock_client.invoke_model_with_response_stream(
+            modelId="anthropic.claude-3-sonnet-20240229-v1:0",
+            body=json.dumps(
+                {
+                    "anthropic_version": "bedrock-2023-05-31",
+                    "max_tokens": 4096,
+                    "system": f"""Here is some relevant information that may help answer the user's upcoming question:
+    
+                    <relevant_information>
+                    {similar_content}
+                    </relevant_information>
+                    
+                    The user's question is:
+                    <question>{question}</question>
+    
+                    Please carefully review the relevant information provided above. 
+    
+                    Your task is to review the provided relevant information and answer the user's question to the best of your ability. 
+                    Aim to use information from the relevant information section to directly address the question asked, and refrain from saying 
+                    things like 'According to the relevant information provided'.
+    
+                    Format your output nicely with sentences that are not too long, in a professional and kind tone. You should prefer lists or 
+                    bullet points when applicable. Begin by thanking the user for their question, and at the end of your answer, say "Thank you for using Ask Dane Street!"
+                    Remember, aim to only use information from the relevant information section in your response, without explicitly referring 
+                    to that section. Return your answer immediately and without preamble. 
+                    </Task>
+                    
+                    <Relevant Information>
+                    {similar_content}
+                    </Relevant Information>""",
+                    "messages": [
+                            {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": message
+                                }
+                            ]
+                        }
+                    ],
+                }
+            ),
+        )
+    
+        # Stream the response
+        all_text = ''
+        stream = response.get('body')
+        if stream:
+            for event in stream:
+                chunk = event.get('chunk')
+                if chunk and json.loads(chunk.get('bytes').decode()):
+                    # check if delta is present
+                    try:
+                        this_text = json.loads(chunk.get('bytes').decode()).get('delta').get('text')
+                        all_text += this_text
+                        yield all_text  # Stream the text back to the UI
+                    except:
+                        pass
+    
+        # Print relevant files
+        output = '\n\nCheck out the following documents for more information:\n'
+        for file, sim in top_docs:
+            output += f"\n{file.replace('.txt', '.pdf')}"
+    
+        yield all_text + output
 
 
 # Create necessary services and collect data