fixed language check and hallu check, added language check to input rails, fixed email sanitize, removed dummy relevance score, refactorings, fixed requirements for installments, refactored model initialization, adapted CLI output

.gitignore

@@ -5,6 +5,7 @@ env/
 .venv
 
 notes.txt
+test.ipynb
 
 # Python
 __pycache__/

app.py

@@ -10,17 +10,30 @@ from rails.output import OutputGuardrails
 import secrets_local
 from helper import ROLE_ASSISTANT, AUTO_ANSWERS
 from rag import retriever
+from dataclasses import dataclass
+from typing import List
 
 # ============================================
 # BACKEND INTEGRATION
 # ============================================
 
-def query_rag_pipeline(user_query: str, model: RAGModel, output_guardRails: OutputGuardrails) -> dict:
+@dataclass
+class Answer(): 
+    answer: str
+    sources: List[str]
+    processing_time: float
+
+def query_rag_pipeline(user_query: str, model: RAGModel, output_guardRails: OutputGuardrails) -> Answer:
     start_time = time.time()
 
     # 1. Input Guardrails
-    if not input_guard.is_valid(user_query):
-        return {"answer": "Invalid query. Please try again.", "sources": []}
+    checked_answer = input_guard.is_valid(user_query)
+    if not checked_answer.accepted:
+        return Answer(
+            answer=checked_answer.reason if checked_answer.reason else "Invalid input. Please try again.", 
+            sources=[], 
+            processing_time=start_time - time.time()
+        )
 
     # 2. Context Retrieval
     retrieved_docs = retriever.search(user_query)
@@ -35,19 +48,20 @@ def query_rag_pipeline(user_query: str, model: RAGModel, output_guardRails: Outp
 
     # 5. Final Answer
     if all(gr_result.passed for gr_result in gr_result.values()): 
-        return  {
-            "answer": response, #output_guardRails.sanitize(response), 
-            "sources": [{"title": str(doc), "relevance": 1.0} for doc in retrieved_docs], # TODO: relevance should be indicated by order in context list
-            "processing_time": end_time - start_time
-        }
+        sources = [{"title": str(output_guardRails.sanitize(doc))} for doc in retrieved_docs]
+        return Answer(
+            answer = output_guardRails.sanitize(response), 
+            sources = sources,
+            processing_time =  end_time - start_time
+        )
     
     else: 
-        issues = [gr_result.issues for gr_result in gr_result.values()]
-        return {
-            "answer": "\n ".join(issues),
-            "sources": [{"title": str(doc), "relevance": 1.0} for doc in retrieved_docs],
-            "processing_time": end_time - start_time
-        }
+        return Answer(
+            answer = output_guardRails.format_guardrail_issues(gr_result),
+            sources = sources,
+            processing_time = end_time - start_time
+        )
+
     
 
 
@@ -90,11 +104,11 @@ def main():
             st.write(message["content"])
             
             # Zeige Quellen wenn vorhanden
-            if message.get("sources"):
-                with st.expander("📚 Verwendete Quellen"):
-                    for source in message["sources"]:
-                        relevance_percent = source["relevance"] * 100
-                        st.write(f"• {source['title']} (Relevanz: {relevance_percent:.0f}%)")
+            if isinstance(message, Answer): 
+                if message.sources:
+                    with st.expander("📚 Verwendete Quellen"):
+                        for source in message.sources:
+                            st.write(f"• {source['title']}")
     
     # Chat-Eingabe
     if prompt := st.chat_input("Stelle eine Frage..."):
@@ -111,21 +125,20 @@ def main():
                 response = query_rag_pipeline(prompt, model, output_guardrails)
 
                 # Antwort anzeigen
-                st.write(response["answer"] if response["answer"] else AUTO_ANSWERS.UNEXPECTED_ERROR)
+                st.write(response.answer if response.answer else AUTO_ANSWERS.UNEXPECTED_ERROR.value)
                 
                 # Antwort in Historie speichern
                 st.session_state.messages.append({
                     "role": ROLE_ASSISTANT,
-                    "content": response["answer"],
-                    "sources": response.get("sources", [])
+                    "content": response.answer,
+                    "sources": response.sources
                 })
                 
                 # Quellen anzeigen wenn vorhanden
-                if response.get("sources"):
+                if response.sources:
                     with st.expander("📚 Verwendete Quellen"):
-                        for source in response["sources"]:
-                            relevance_percent = source["relevance"] * 100
-                            st.write(f"• {source['title']} (Relevanz: {relevance_percent:.0f}%)")
+                        for source in response.sources:
+                            st.write(f"• {source['title']}")
             
         
 

helper.py

@@ -1,12 +1,63 @@
 from enum import Enum
 import re 
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline, AutoTokenizer
+from functools import lru_cache
+
+
 
 MODEL_NAME = 'google/gemma-2-2b-it'
-SIMILARITY_MODEL_NAME = 'all-MiniLM-L6-v2'
-HALLUCINATION_MODEL_NAME="vectara/hallucination_evaluation_model"
 ROLE_ASSISTANT ="assistant"
 EMAIL_PATTERN = re.compile(r"[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+")
 
+@lru_cache(maxsize=1)
+def get_similarity_model():
+    return SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2', device="cpu")
+
+@lru_cache(maxsize=1)
+def get_toxicity_model(): 
+    return pipeline(
+            "text-classification",
+            model="unitary/toxic-bert",
+            top_k=None
+        )
+
+@lru_cache(maxsize=1)
+def get_hallucination_model(): 
+    return pipeline(
+            "text-classification",
+            model='vectara/hallucination_evaluation_model',
+            tokenizer=AutoTokenizer.from_pretrained('google/flan-t5-base'),
+            trust_remote_code=True
+        )
+
+def check_toxicity(text: str): 
+    """ Toxicity check with toxicity check pipeline
+        Args: 
+            response (str): response generated by model
+        Returns: 
+
+    """
+    try:
+        # output sequence can be little longer due to tokens != sequence length, this makes sure that text is cropped to model input max size
+        if len(text) > 512: 
+            text = text[:500]
+
+        result = get_toxicity_model()(text)[0]
+        # print(result)
+        if result:
+            toxicity_score = next((s['score'] for s in result if s['score'] > 0.85), sum(s['score'] for s in result) / len(result))
+            # print(toxicity_score)
+            toxicity_passed = True if toxicity_score < 0.6 else False
+            text = AUTO_ANSWERS.LANGUAGE_INAPPROPRIATE.value if not toxicity_passed else ""
+            return toxicity_passed, toxicity_score, text
+        else: 
+            raise ValueError("Could not yield result.")
+    except Exception as e:
+        print(f"Error while checking language: {e}")
+        return False, 1.0, AUTO_ANSWERS.UNEXPECTED_ERROR.value
+    
+
 class AUTO_ANSWERS(Enum): 
     COULD_NOT_GENERATE = "Could not generate an answer."
     REQUEST_TIMED_OUT = "The request timed out."
@@ -17,4 +68,6 @@ class AUTO_ANSWERS(Enum):
     HALLUCINATION_CHECK_FAILED = "Hallucination check failed."
     NOT_SUPPORTED_BY_CONTEXT = "Many facts not supported by context."
     INCONSISTENT_WITH_CONTEXT = "Inconsistent with context."
-    CORRECTNESS_CHECK_FAILED = "Correctness check failed."
+    CORRECTNESS_CHECK_FAILED = "Correctness check failed."
+    LANGUAGE_INAPPROPRIATE = "Inappropriate language detected!"
+    REPHRASE_SENTENCE = "Try to rephrase your request."

model/model.py

@@ -60,9 +60,7 @@ class RAGModel:
                 temperature=0.7,         # Adjust creativity (higher = more random)
                 top_p=0.9,               # Nucleus sampling parameter
             )
-            
-            print(f"RESPONSE {response}")
-        
+                    
             if response: 
                 return response.choices[0].message.content
             else:

rag/retriever.py

@@ -1,6 +1,7 @@
 import chromadb
 from sentence_transformers import SentenceTransformer
 import sqlite3
+from helper import get_similarity_model
 
 def search(query: str, top_k: int = 5):
     """
@@ -14,10 +15,10 @@ def search(query: str, top_k: int = 5):
         conn.close()
         return [{"name": student[0]} for student in students]
 
-    # Initialize the embedding model
-    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    # get SentenceTransformer model
+    model = get_similarity_model()
 
-    # Create the query embedding
+    # Create the query embedding - model from helper instantiated only once 
     query_embedding = model.encode([query])
 
     # Initialize ChromaDB client and get the collection

rails/input.py

@@ -1,12 +1,20 @@
 import re
+from helper import check_toxicity, AUTO_ANSWERS
+from typing import Optional
+from dataclasses import dataclass
 
-def is_valid(query: str) -> bool:
+@dataclass
+class CheckedInput(): 
+    accepted: bool
+    reason: Optional[str]
+
+def is_valid(query: str) -> CheckedInput:
     """
     Validates the user's query.
     """
     # Check for query length
     if len(query) > 500:
-        return False
+        return CheckedInput(False, get_output("Input too long."))
 
     # Check for SQL injection patterns
     sql_injection_patterns = [
@@ -15,9 +23,16 @@ def is_valid(query: str) -> bool:
     ]
     for pattern in sql_injection_patterns:
         if re.search(pattern, query, re.IGNORECASE):
-            return False
+            return CheckedInput(False, get_output("Invalid input."))
+        
+    t_passed, _ , text = check_toxicity(query)
+    if not t_passed: 
+        return CheckedInput(False, get_output(text))
+
+    return CheckedInput(True, None)
 
-    return True
+def get_output(reason: str) -> str: 
+    return "\n".join([reason, AUTO_ANSWERS.REPHRASE_SENTENCE.value])
 
 if __name__ == '__main__':
     # Example usage

rails/output.py

@@ -1,21 +1,16 @@
 from sentence_transformers import SentenceTransformer
 from typing import Dict
-from helper import SIMILARITY_MODEL_NAME, HALLUCINATION_MODEL_NAME, EMAIL_PATTERN
-from transformers import pipeline
+from helper import EMAIL_PATTERN, get_similarity_model, AUTO_ANSWERS, get_toxicity_model, get_hallucination_model, check_toxicity
+from transformers import pipeline, AutoTokenizer
 from sklearn.metrics.pairwise import cosine_similarity
-#from guardrails import Guard
-#from guardrails.hub import DetectPII, ToxicLanguage
-#from guardrails.validators import ValidLength
-from helper import AUTO_ANSWERS
 from typing import Dict, List
 from dataclasses import dataclass
 from enum import Enum
-import torch
 import re
 
 @dataclass
 class GuardrailResult:
-    """Ergebnis der Guardrail-Prüfung"""
+    """Guardrail Result Dataclass"""
     passed: bool
     confidence: float
     issues: List[str]
@@ -25,10 +20,10 @@ class GuardrailType(Enum):
     RELEVANCE = "relevance"
     HALLUCINATION = "hallucination"
     CORRECTNESS = "correctness"
-    #TOXICITY = "toxicity"
+    TOXICITY = "toxicity"
 
 class OutputGuardrails:
-    """Umfassende Output Guardrails für RAG Pipeline"""
+    """Guardrails for LLM output validation"""
     
     def __init__(self):
 
@@ -43,27 +38,22 @@ class OutputGuardrails:
             re.compile(r'\b\d+%\b'),   # %
             re.compile(r'\b\d+\s*(Euro|Dollar|meter|kilometer|km|kg|tons)\b', re.IGNORECASE),  
         ]
-
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.sentence_model = SentenceTransformer(SIMILARITY_MODEL_NAME, device=device)
-        #self.sentence_model = self.sentence_model.to(device)
         
-        # Hallucination Detection Model
         try:
-            self.hallucination_model = pipeline(
-                "text-classification",
-                model=HALLUCINATION_MODEL_NAME,
-                trust_remote_code=True
-            )
+            self.sentence_model = get_similarity_model()
+        except Exception as e:
+            print(f"Could not retrieve sentence model for similarity checks: {e}.")
+            
+        try:
+            self.hallucination_model = get_hallucination_model()
         except Exception as e:
             print(f"Could not retrieve hallucination check pipeline: {e}.")
-            self.hallucination_model = None
-        
-        '''  # setup guard for guardrails - toxicity check - disabled for now, work in progress
-        self.guard = Guard().use_many(
-            ToxicLanguage(threshold=0.8),
-            ValidLength(min=10, max=1000)
-        )'''
+
+        try:
+            self.toxicity_classifier = get_toxicity_model()
+        except Exception as e:
+            print(f"Could not toxicity check pipeline: {e}.")
+       
 
 
     def check_query_relevance(self, 
@@ -98,7 +88,7 @@ class OutputGuardrails:
             if relevance_score < 0.3:
                 issues.append(AUTO_ANSWERS.ANSWER_NOT_RELEVANT_TO_QUERY.value)
             if context_relevance < 0.4:
-                issues.append(AUTO_ANSWERS.ANSWER_NOT_RELEVANT_TO_CONTEXT)
+                issues.append(AUTO_ANSWERS.ANSWER_NOT_RELEVANT_TO_CONTEXT.value)
             
             passed = combined_relevance >= 0.5 and len(issues) == 0
             
@@ -167,18 +157,25 @@ class OutputGuardrails:
             
             # Check if model facts are consistent with context
             if self.hallucination_model:
-                input_text = f"Context: {context_text}\nClaim: {response}"
-                result = self.hallucination_model(input_text)
-                hallucination_score = next(
-                    (score['score'] for score in result if score['label'] == 'HALLUCINATION'), 
-                    0.0
-                )
+                
+                prompt = f"<pad> Determine if the hypothesis is true given the premise?\n\nPremise: {context}\n\nHypothesis: {response}"
+
+                classifier = pipeline(
+                            "text-classification",
+                            model='vectara/hallucination_evaluation_model',
+                            tokenizer=AutoTokenizer.from_pretrained('google/flan-t5-base'),
+                            trust_remote_code=True
+                        )
+                full_scores = classifier(prompt, top_k=None) # List[List[Dict[str, float]]]
+                hallucinated_entry = next((item for item in full_scores if item['label'] == 'hallucinated'), None)
+                hallucination_score = hallucinated_entry['score']
+                
             else:
                 # Do semantic similarity check if hallucination model pipeline is unavail
                 context_embedding = self.sentence_model.encode([context_text])
                 response_embedding = self.sentence_model.encode([response])
                 similarity = cosine_similarity(context_embedding, response_embedding)[0][0]
-                hallucination_score = 1.0 - similarity  # Invertiert für Hallucination Score
+                hallucination_score = 1.0 - similarity  
             
             # check for known hallucination patterns                  
             pattern_score = 0
@@ -192,7 +189,7 @@ class OutputGuardrails:
             
             # Gesamtscore
             total_hallucination_score = min(
-                hallucination_score + pattern_score + contradiction_score, 1.0
+                hallucination_score*0.4 + pattern_score*0.3 + contradiction_score*0.3, 1.0
             )
             
             if total_hallucination_score > 0.7:
@@ -323,12 +320,9 @@ class OutputGuardrails:
             else:
                 support_ratio = 1.0 # no checkable facts 
             
-            # consistency with context
+     
             consistency_score = self._check_consistency(response, context_text)
-            
             completeness_score = self._check_completeness(response, context_text)
-            
-            # Overall, might want to adapt weights 
             correctness_score = (
                 support_ratio * 0.5 + 
                 consistency_score * 0.3 + 
@@ -358,25 +352,27 @@ class OutputGuardrails:
             print(f"{AUTO_ANSWERS.CORRECTNESS_CHECK_FAILED.value} {e}")
             return GuardrailResult(False, 0.0, [AUTO_ANSWERS.CORRECTNESS_CHECK_FAILED.value], {})
     
-    '''
-    def toxicity_check(self, response: str): 
-        """ Toxicity Check with Guardrails Library
+    
+    def check_toxicity(self, response: str): 
+        """ Toxicity check with toxicity check pipeline
             Args: 
                 response (str): response generated by model
+            Returns: 
+
         """
-        try:
-            guard_result = self.guard.validate(response)
-            toxicity_passed = guard_result.validation_passed
-            toxicity_issues = [str(error) for error in guard_result.reask_messages] if hasattr(guard_result, 'reask_messages') else []
-            return toxicity_passed, toxicity_issues
-        except Exception as e:
-            print(f"Guardrails toxicity validation failed: {e}")
-            toxicity_passed = False
-            toxicity_issues = []
-    '''
+        print("CHECKING OUTPUT LANGUAGE: ")
+        toxicity_passed, toxicity_score, text = check_toxicity(response)
 
+        return GuardrailResult(
+            passed=toxicity_passed,
+            confidence=1 - toxicity_score,
+            issues=[text] if text else [],
+            metrics={"toxicity_passed": toxicity_passed}
+        )
+            
     def sanitize(self, response: str) -> str: 
-        return EMAIL_PATTERN.sub(response, '[REDACTED_EMAIL]')
+        EMAIL_PATTERN = re.compile(r"[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+")
+        return EMAIL_PATTERN.sub('[REDACTED_EMAIL]', response)
 
     def check(self, 
                                  query: str, 
@@ -395,28 +391,57 @@ class OutputGuardrails:
             self.check_query_relevance(query, response, context),
             self.check_hallucination(response, context),
             self.check_correctness(response, context),
+            self.check_toxicity(response)
         ]
         
-        relevance_result, hallucination_result, correctness_result = tasks
+        relevance_result, hallucination_result, correctness_result, toxicity_result = tasks
         
         results[GuardrailType.RELEVANCE] = relevance_result
         results[GuardrailType.HALLUCINATION] = hallucination_result
         results[GuardrailType.CORRECTNESS] = correctness_result
-        
-        '''
-        result = self.toxicity_check(response)
-        if result: 
-            toxicity_passed, toxicity_issues = result 
-            results[GuardrailType.TOXICITY] = GuardrailResult(
-                passed=toxicity_passed,
-                confidence=1.0 if toxicity_passed else 0.0,
-                issues=toxicity_issues,
-                metrics={"toxicity_passed": toxicity_passed}
-            )
+        results[GuardrailType.TOXICITY] = toxicity_result
+
 
-        else: 
-            # Todo: discuss whether how relevant this could be in our application
-            print("WARNING: Output could be toxic.")
-        '''
+        print("Result for query: ", query)
+        print("WITH CONTEXT: ", context)
+        print("Model output: ", response)
+        for guardrail_type, result in results.items():
+            print(f"\n Guardrail: {guardrail_type.value.upper()}")
+            print(f"   Passed:     {result.passed}")
+            print(f"   Confidence: {result.confidence:.3f}")
+            
+            if result.issues:
+                print(f"   Issues:     {', '.join(result.issues)}")
+            else:
+                print(f"   Issues:     None")
             
-        return results
+            print("   Metrics:")
+            for metric, value in result.metrics.items():
+                print(f"     - {metric}: {value:.4f}")
+            
+        return results
+    
+
+    def format_guardrail_issues(self, gr_result: dict) -> str:
+        """Formats the issues text to readable output
+        Parametes: 
+            - gr_result (dict): The dictionary containing the different results from the output guardrail check
+        Returns: 
+            Formatted string. 
+        """
+        
+        lines = ["Output Issues:\n"]
+
+        for guardrail_type, result in gr_result.items():
+            issues = result.issues
+            title = f"- {guardrail_type.value.capitalize()} issues:"
+            if issues:
+                lines.append(title)
+                for issue in issues:
+                    lines.append(f"  • {issue}")
+                if guardrail_type == GuardrailType.TOXICITY and not result.passed: 
+                    lines.append(f"  • {AUTO_ANSWERS.REPHRASE_SENTENCE.value}")
+            else:
+                lines.append(f"{title} None")
+        
+        return "\n".join(lines)

requirements.txt

@@ -3,8 +3,9 @@ python-dotenv==1.0.1
 transformers
 sentence-transformers
 scikit-learn
-Faker==15.3.4
+faker
 chromadb==1.0.21
 sentence-transformers==5.1.0
 numpy==1.26.4
 huggingface-hub==0.34.4
+chromadb