Update app.py

app.py

@@ -16,7 +16,16 @@ from langchain.llms import HuggingFacePipeline
 import gradio as gr
 import re
 from bs4 import BeautifulSoup
+from guardrails.validators import Validator, register_validator, ValidationResult, FailResult, PassResult
+from presidio_analyzer import AnalyzerEngine
+from presidio_analyzer.nlp_engine import SpacyNlpEngine, NlpEngineProvider
+from better_profanity import profanity 
+from presidio_analyzer import PatternRecognizer, Pattern
 import inflection
+from guardrails import Guard
+import warnings
+# Suppress all warnings
+warnings.filterwarnings("ignore")
 
 
 """
@@ -119,10 +128,76 @@ def get_confidence_score(question):
     return min(1.0, round(max_score, 2))  # Normalize to 0-1 scale
 
 
+## GuardRail validators
+
+# Define NLP Configuration with lang_code
+nlp_configuration = {
+    "nlp_engine_name": "spacy",
+    "models": [{"lang_code": "en", "model_name": "en_core_web_lg"}],  # Specify language
+}
+
+# Define SSN Pattern
+ssn_regex = r"\b\d{3}-\d{2}-\d{4}\b"  # Matches US SSN format (123-45-6789)
+ssn_pattern = Pattern(name="SSN Pattern", regex=ssn_regex, score=0.85)  # Score between 0-1
+
+# Create Custom SSN Recognizer
+ssn_recognizer = PatternRecognizer(supported_entity="SSN", patterns=[ssn_pattern])
+
+analyzer = AnalyzerEngine()
+analyzer.registry.add_recognizer(ssn_recognizer)
+
+@register_validator(name="custom_pii_detector", data_type="string")
+class CustomPIIDetector(Validator):
+
+    def validate(self, value, metadata={}) -> ValidationResult:
+        # Analyze text for PII
+        results = analyzer.analyze(text=value, entities=["PHONE_NUMBER", "EMAIL_ADDRESS", "CREDIT_CARD", "SSN"], language="en")
+
+        if results:
+            detected_entities = ", ".join(set([res.entity_type for res in results]))
+            return FailResult(
+                error_message=f"Query contains PII: {detected_entities}."
+            )
+
+        return PassResult()
+
+# Custom Profanity Detector using better-profanity
+@register_validator(name="custom_profanity_detector", data_type="string")
+class CustomProfanityDetector(Validator):
+    def validate(self, value, metadata={}) -> ValidationResult:
+        if profanity.contains_profanity(value):
+            return FailResult(
+                error_message="Query contains profanity."
+            )
+        return PassResult()
+
+# Custom Relevance Validator for Finance and Apple-related Queries
+@register_validator(name="custom_relevance_detector", data_type="string")
+class CustomRelevanceDetector(Validator):
+    def validate(self, value, metadata={}) -> ValidationResult:
+        finance_keywords = {"revenue", "profit", "expenses", "balance sheet", "earnings", "financial", "investment", "dividends", "assets", "liabilities", "cash flow", "loss","turnover"}
+        apple_keywords = {"apple", "iphone", "macbook", "tim cook", "apple inc", "ios", "mac", "ipad"}
+        
+        text_lower = value.lower()
+        
+        # Check if any finance-related or Apple-related keyword appears in the query
+        if not any(keyword in text_lower for keyword in (finance_keywords | apple_keywords)):  # Use set union
+            return FailResult(
+                error_message="Query is not related to finance or Apple."
+            )
+        
+        return PassResult()
+
+guard = Guard().use(CustomPIIDetector).use(CustomProfanityDetector).use(CustomRelevanceDetector)
+
 """### 📌 Step 10: Integrate with Gradio UI"""
 
 # Define Chatbot Function
 def chat_with_rag(message, history):
+    try:
+        res = guard.validate(message)
+    except Exception as e:
+        return f"❌ Guardrail {str(e)}"
     try:
         response = conversation_chain.invoke(message)
         confidence_score = get_confidence_score(message)
@@ -155,19 +230,18 @@ print(f"📌 **Answer:** {to_camel_case(output)}\n\n**Confidence Score:** {confi
 # Create Gradio Chatbot UI with Auto-Clearing Input
 demo = gr.ChatInterface(
     fn=chat_with_rag,  # Function to generate responses
-    title="📊 Financial RAG Chatbot",
-    description="Ask questions about financial reports and get AI-powered answers!",
-    theme="soft",  # Aesthetic theme
+    title="📊 Financial Basic RAG Chatbot",
+    description="Ask questions about Apple's financial reports and get AI-powered answers!",
+    theme="soft", 
     examples=[
         ["What are the biggest challenges for Apple?"],
         ["What was Apple's total revenue in 2024?"],
         ["What are the biggest financial risks for Apple?"],
-        ["What is the capital of France?"]
+        ["What is the capital of France?"],
     ],
-    submit_btn="Ask",  # Customize the submit button text
-    stop_btn=None,  # Remove unnecessary stop button
+    submit_btn="Ask",  
+    stop_btn=None,  
 )
 
-
 if __name__ == "__main__":
     demo.launch()