Update prompts.py

prompts.py

@@ -1,48 +1,39 @@
-# prompts.py
 from langchain.prompts import PromptTemplate
 
+# Define a list of harmful or inappropriate topics to block
+HARMFUL_TOPICS = ["kill", "suicide", "harm", "death", "murder", "violence"]
+
 classification_prompt_str = """
 You are a helpful assistant that classifies user questions into three categories:
 1) "Wellness" if the question involves health, nutrition, fitness, mental well-being, self-care, or research related to these.
 2) "Brand" if the question specifically pertains to 'DailyWellnessAI'—its mission, disclaimers, features, policies, etc.
-3) "OutOfScope" if it does not fall into the above two categories.
-
+3) "OutOfScope" if it does not fall into the above two categories or if the question involves harmful topics.
 **Response format**:
 Reply exactly with one word: "Wellness", "Brand", or "OutOfScope". Do not provide any additional explanation.
-
 Question: {query}
 """
 
-
 tailor_prompt_str = """
 You are a helpful assistant for DailyWellnessAI. Your role is to simplify complex ideas and offer actionable, user-friendly advice that aligns with our mission to enhance daily wellness through AI.
-
 Here's the response to tailor:
 {response}
-
 Tailor it to ensure:
 - Simplicity and clarity.
 - Practicality, with actionable recommendations where appropriate.
 - Alignment with DailyWellnessAI's mission of simplifying daily wellness through AI.
-
 Provide the revised, concise response below:
 """
 
-
 cleaner_prompt_str = """
 You are a helpful AI. Below, you have two sources of information:
-
 1) CSV (Knowledge Base) Answer:
 {kb_answer}
-
 2) Web Search Result:
 {web_answer}
-
 Your task:
 - Combine and synthesize the two into a single, clear, cohesive answer.
 - Avoid duplication or irrelevant details.
 - Present the final response in straightforward, user-friendly language.
-
 Do not repeat content verbatim. Merge the information meaningfully and provide your synthesized answer below:
 """
 
@@ -54,11 +45,8 @@ It can help promote mindfulness, relaxation, balance, or focus, all of which con
 For more wellness-related questions or to learn more about DailyWellnessAI, feel free to ask—I’m here to support your wellness journey!
 """
 
-# Define the PromptTemplate
+# Define the PromptTemplate objects:
 
-
-
-# Now we define the PromptTemplate objects:
 classification_prompt = PromptTemplate(
     template=classification_prompt_str,
     input_variables=["query"]
@@ -78,3 +66,20 @@ refusal_prompt = PromptTemplate(
     template=refusal_prompt_str,
     input_variables=["topic"]
 )
+
+# Define function to check if the query contains harmful topics
+def is_harmful(query: str) -> bool:
+    # Check if the query contains any harmful topics
+    for topic in HARMFUL_TOPICS:
+        if topic.lower() in query.lower():
+            return True
+    return False
+
+# Now modify the classification logic:
+def classify_query(query: str) -> str:
+    if is_harmful(query):
+        return "OutOfScope"
+    # Use the classification prompt template to classify the query
+    classification = classification_prompt.invoke({"query": query}).get("text", "").strip()
+    return classification
+