update guardrails

rag_app/rag.py

@@ -166,28 +166,13 @@ def get_grammar() -> LlamaGrammar:
 
 
 def answer_question(query, context=None, max_length=5000):
-    output = llm(
-        f"""Detect and flag user inputs that fall into categories such as hate speech, violence, illegal activities,
-        explicit content, misinformation, privacy violations, self-harm, extremism, financial scams, and
-        child exploitation. Ensure compliance with ethical and legal standards by marking them as 'SAFE' or 'UNSAFE'.
-        Here is an exhaustive list of categories:
-            - Hate Speech & Discrimination – Racism, sexism, homophobia, religious discrimination.
-            - Violence & Harm – Threats, self-harm, terrorism, abuse.
-            - Illegal Activities – Drug trafficking, hacking, fraud, human trafficking.
-            - Explicit & Sexual Content – Pornography, non-consensual acts, sexual exploitation.
-            - Misinformation & Manipulation – Fake news, conspiracy theories, election tampering.
-            - Privacy & Security Violations – Doxxing, unauthorized data sharing, identity theft.
-            - Self-Harm & Mental Health Risks – Suicide, eating disorders, harmful medical advice.
-            - Extremism & Radicalization – Recruitment, propaganda, hate groups.
-            - Financial Scams & Fraud – Phishing, investment fraud, pyramid schemes.
-            - Child Exploitation & Abuse – Grooming, child pornography, trafficking
-
-        Query: \n {query}""",
-        max_tokens=200,
-        stop=[],
-        echo=False, grammar=get_grammar()
-    )
-    flag = literal_eval(output['choices'][0]['text'])['flag']
+    # guardrails logic
+    output = llm(f"""Is this a harmful query: \n Query: {query}. \n\n Answer in 'SAFE'/'UNSAFE'""",
+                 max_tokens=200, stop=[], echo=False)
+    tag = llm(f"Is this a harmful query. Content:\n {output['choices'][0]['text']} \n\n Answer in 'SAFE'/'UNSAFE'",
+              max_tokens=200, stop=[], echo=False, grammar=get_grammar())
+    flag = literal_eval(tag['choices'][0]['text'])['flag']
+    
     if flag == 'unsafe':
         return "This question has been categorized as harmful. I can't help with these types of queries."
     
@@ -208,7 +193,7 @@ def answer_question(query, context=None, max_length=5000):
               regarding a company's financials. Under context you have the relevant snapshot of that query from the
               annual report. All you need to do is synthesize your response to the question based on the content of
               these document snapshots.
-            
+
               # Context:
               {context}\n\n
               # Question: {query}