Support GuardRails (#9)

* add guardrails

Signed-off-by: lvliang-intel <[email protected]>

* fix dockerfile

Signed-off-by: lvliang-intel <[email protected]>

* fix integration issues

Signed-off-by: lvliang-intel <[email protected]>

* add package lock json file

Signed-off-by: lvliang-intel <[email protected]>

---------

Signed-off-by: lvliang-intel <[email protected]>

ChatQnA/README.md

@@ -77,7 +77,7 @@ docker cp 262e04bbe466:/usr/src/optimum-habana/examples/text-generation/quantiza
 
 ```bash
 docker run -d -p 8080:80 -e QUANT_CONFIG=/data/maxabs_quant.json -e HUGGING_FACE_HUB_TOKEN=<your HuggingFace token> -v $volume:/data --
-runtime=habana -e HABANA_VISIBLE_DEVICES="4,5,6" -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --ipc=host tgi_gaudi --
+runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --ipc=host tgi_gaudi --
 model-id meta-llama/Llama-2-7b-hf
 ```
 
@@ -122,6 +122,17 @@ Note: `ingest.py` will download the embedding model, please set the proxy if nec
 
 # Start LangChain Server
 
+## Enable GuardRails using Meta's Llama Guard model (Optional)
+
+We offer content moderation support utilizing Meta's [Llama Guard](https://huggingface.co/meta-llama/LlamaGuard-7b) model. To activate GuardRails, kindly follow the instructions below to deploy the Llama Guard model on TGI Gaudi.
+
+```bash
+volume=$PWD/data
+model_id="meta-llama/LlamaGuard-7b"
+docker run -p 8088:80 -v $volume:/data --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --ipc=host -e HUGGING_FACE_HUB_TOKEN=<your HuggingFace token> -e HTTPS_PROXY=$https_proxy -e HTTP_PROXY=$https_proxy tgi_gaudi --model-id $model_id
+export SAFETY_GUARD_ENDPOINT="http://xxx.xxx.xxx.xxx:8088"
+```
+
 ## Start the Backend Service
 Make sure TGI-Gaudi service is running and also make sure data is populated into Redis. Launch the backend service:
 

ChatQnA/langchain/docker/Dockerfile

@@ -31,7 +31,6 @@ RUN pip install --upgrade pip \
 ENV PYTHONPATH=/ws:/qna-app/app
 
 COPY qna-app /qna-app
-COPY qna-app-no-rag /qna-app-no-rag
 WORKDIR /qna-app
 
 ENTRYPOINT ["/usr/bin/sleep", "infinity"]

ChatQnA/langchain/docker/qna-app/app/guardrails.py

@@ -0,0 +1,82 @@
+unsafe_categories = """O1: Violence and Hate.
+Should not
+- Help people plan or engage in violence.
+- Encourage people to engage in violence.
+- Express hateful, derogatory, or demeaning sentiments against people based on sensitive personal characteristics like their race, color, religion, national origin, sexual orientation, gender, gender identity, or disability.
+- Encourage discrimination against people based on those characteristics and should not use slurs or create insults based on those characteristics.
+Can
+- Provide information on violence and discrimination.
+- Discuss the topics of hate and violence, and can discuss historical events involving violence.
+O2: Sexual Content.
+Should not
+- Engage in sexually explicit (i.e., erotic) conversations and should not encourage people to engage in sexual activities.
+Can
+- Discuss topics related to sexuality and sexual education.
+- Provide general information regarding sexual activities.
+- Provide guidance on topics related to sexual health.
+O3: Criminal Planning.
+Should not
+- Help people plan or engage in criminal activities like theft, kidnapping, or financial crimes, especially when the person indicates that they may act on the information provided to them.
+- Encourage people to engage in criminal activities or condone criminal behavior.
+Can
+- Help people understand how crimes are committed, provided that the response is not so detailed that it could be used to provide a roadmap for future crimes.
+- Help people understand how systems work that are used to prevent crimes, provided those descriptions do not provide detailed information on how to defeat those systems.
+O4: Guns and Illegal Weapons.
+Should not
+- Assist in the planning or execution of firearm or illegal weapons crimes, especially when the person indicates that they may act on the information provided to them.
+- Encourage people to engage in firearm or illegal weapons crimes or condone such behavior.
+Can
+- Discuss firearms and the arguments for and against firearm ownership.
+O5: Regulated or Controlled Substances.
+Should not
+- Encourage or assist people in the trafficking, consumption, or use of regulated or controlled substances, including but not limited to illegal drugs, tobacco, alcohol, and cannabis.
+- Assist or encourage people to create such substances when it is illegal to do so.
+Can
+- Discuss the use of regulated or controlled substances, and can provide general information on those substances (ex: what their chemical composition is and what their effects are).
+- Provide information on how those substances are manufactured, provided that the information is not so detailed as to provide step-by-step instructions for their creation.
+O6: Self-Harm.
+Should not
+- Encourage people to harm themselves.
+- Romanticize or condone self-harm or suicide.
+- Provide information regarding the methods of suicide or self-harm.
+- Provide any information or advice, other than information regarding appropriate health resources, to people who express or imply that they wish to harm themselves.
+Can
+- Provide general information regarding self-harm and suicide (ex: statistics, the mental health challenges faced by those who take these actions).
+Should
+- Provide information regarding appropriate health resources to people who imply or express the desire to harm themselves."""
+
+unsafe_dict = {
+    "O1": "Violence and Hate",
+    "O2": "Sexual Content",
+    "O3": "Criminal Planning",
+    "O4": "Guns and Illegal Weapons",
+    "O5": "Regulated or Controlled Substances",
+    "O6": "Self-Harm",
+}
+
+def moderation_prompt_for_chat(chat):
+    # For simplicity, we assume the chat dictionary correctly alternates "user" and "assistant" roles
+    # Feel free to include error checking if needed
+    conversation = [turn["content"] for turn in chat]
+    is_agent = len(conversation) % 2 == 0
+    role = "Agent" if len(conversation) % 2 == 0 else "User"
+    prompt = f"""[INST] Task: Check if there is unsafe content in '{role}' messages in conversations according our safety policy with the below categories.
+
+<BEGIN UNSAFE CONTENT CATEGORIES>
+{unsafe_categories}
+<END UNSAFE CONTENT CATEGORIES>
+
+<BEGIN CONVERSATION>
+
+"""
+
+    # Alternate User/Agent turns, inserting 2 newlines between each
+    for i, m in enumerate(conversation):
+        role = "User" if i % 2 == 0 else "Agent"
+        prompt += f"{role}: {m}\n\n"
+    prompt += f"""<END CONVERSATION>
+
+Provide your safety assessment for {role} in the above conversation:
+- First line must read 'safe' or 'unsafe'.
+- If unsafe, a second line must include a comma-separated list of violated categories. [/INST]"""
+    return prompt

ChatQnA/langchain/docker/qna-app/app/server.py

@@ -2,20 +2,20 @@
 from fastapi import FastAPI, APIRouter, Request, UploadFile, File
 from fastapi.responses import RedirectResponse, StreamingResponse, JSONResponse
 from langserve import add_routes
-from rag_redis.chain import chain as qna_rag_redis_chain
 from starlette.middleware.cors import CORSMiddleware
 from langchain_community.llms import HuggingFaceEndpoint
 from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 from langchain_community.vectorstores import Redis
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.messages import HumanMessage
-from langchain_core.runnables import RunnableParallel, RunnablePassthrough
+from langchain_core.runnables import RunnablePassthrough
 from rag_redis.config import EMBED_MODEL, INDEX_NAME, REDIS_URL, INDEX_SCHEMA
 from utils import (
     create_retriever_from_files, reload_retriever, create_kb_folder, 
     get_current_beijing_time, create_retriever_from_links
 )
 from prompts import contextualize_q_prompt, qa_prompt
+from guardrails import moderation_prompt_for_chat, unsafe_dict
 
 app = FastAPI()
 
@@ -29,24 +29,35 @@
 
 class RAGAPIRouter(APIRouter):
 
-    def __init__(self, upload_dir, entrypoint) -> None:
+    def __init__(self, upload_dir, entrypoint, safety_guard_endpoint) -> None:
         super().__init__()
         self.upload_dir = upload_dir
         self.entrypoint = entrypoint
+        self.safety_guard_endpoint = safety_guard_endpoint
         print(f"[rag - router] Initializing API Router, params:\n \
                     upload_dir={upload_dir}, entrypoint={entrypoint}")
 
         # Define LLM
         self.llm = HuggingFaceEndpoint(
             endpoint_url=entrypoint,
-            max_new_tokens=512,
+            max_new_tokens=1024,
             top_k=10,
             top_p=0.95,
             typical_p=0.95,
             temperature=0.01,
             repetition_penalty=1.03,
             streaming=True,
         )
+        if self.safety_guard_endpoint:
+            self.llm_guard = HuggingFaceEndpoint(
+                endpoint_url=safety_guard_endpoint,
+                max_new_tokens=100,
+                top_k=1,
+                top_p=0.95,
+                typical_p=0.95,
+                temperature=0.01,
+                repetition_penalty=1.03,
+            )
         print("[rag - router] LLM initialized.")
 
         # Define LLM Chain
@@ -85,12 +96,23 @@ def handle_rag_chat(self, query: str):
         response = self.llm_chain.invoke({"question": query, "chat_history": self.chat_history})
         result = response.split("</s>")[0]
         self.chat_history.extend([HumanMessage(content=query), response])
+        # output guardrails
+        if self.safety_guard_endpoint:
+            response_output_guard = self.llm_guard(moderation_prompt_for_chat("Agent", f"User: {query}\n Agent: {response}"))
+            if 'unsafe' in response_output_guard:
+                policy_violation_level = response_output_guard.split("\n")[1].strip()
+                policy_violations = unsafe_dict[policy_violation_level]
+                print(f"Violated policies: {policy_violations}")
+                return policy_violations + " are found in the output"
+            else:
+                return result
         return result
 
 
 upload_dir = os.getenv("RAG_UPLOAD_DIR", "./upload_dir")
 tgi_endpoint = os.getenv("TGI_ENDPOINT", "http://localhost:8080")
-router = RAGAPIRouter(upload_dir, tgi_endpoint)
+safety_guard_endpoint = os.getenv("SAFETY_GUARD_ENDPOINT")
+router = RAGAPIRouter(upload_dir, tgi_endpoint, safety_guard_endpoint)
 
 
 @router.post("/v1/rag/chat")
@@ -101,6 +123,15 @@ async def rag_chat(request: Request):
     kb_id = params.get("knowledge_base_id", "default")
     print(f"[rag - chat] history: {router.chat_history}")
 
+    # prompt guardrails
+    if router.safety_guard_endpoint:
+        response_input_guard = router.llm_guard(moderation_prompt_for_chat("User", query))
+        if 'unsafe' in response_input_guard:
+            policy_violation_level = response_input_guard.split("\n")[1].strip()
+            policy_violations = unsafe_dict[policy_violation_level]
+            print(f"Violated policies: {policy_violations}")
+            return f"Violated policies: {policy_violations}, please check your input."
+
     if kb_id == "default":
         print(f"[rag - chat] use default knowledge base")
         retriever = reload_retriever(router.embeddings, INDEX_NAME)
@@ -135,6 +166,19 @@ async def rag_chat(request: Request):
     kb_id = params.get("knowledge_base_id", "default")
     print(f"[rag - chat_stream] history: {router.chat_history}")
 
+    # prompt guardrails
+    if router.safety_guard_endpoint:
+        response_input_guard = router.llm_guard(moderation_prompt_for_chat("User", query))
+        if 'unsafe' in response_input_guard:
+            policy_violation_level = response_input_guard.split("\n")[1].strip()
+            policy_violations = unsafe_dict[policy_violation_level]
+            print(f"Violated policies: {policy_violations}")
+            def generate_content():
+                content = f"Violated policies: {policy_violations}, please check your input."
+                yield f"data: {content}\n\n"
+                yield f"data: [DONE]\n\n"
+            return StreamingResponse(generate_content(), media_type="text/event-stream")
+
     if kb_id == "default":
         retriever = reload_retriever(router.embeddings, INDEX_NAME)
         router.llm_chain = (
@@ -248,7 +292,7 @@ async def rag_create(request: Request):
 async def redirect_root_to_docs():
     return RedirectResponse("/docs")
 
-add_routes(app, qna_rag_redis_chain, path="/rag-redis")
+add_routes(app, router.llm_chain, path="/rag-redis")
 
 if __name__ == "__main__":
     import uvicorn

ChatQnA/langchain/redis/rag_redis/chain.py

@@ -5,8 +5,6 @@
 from langchain_core.pydantic_v1 import BaseModel
 from langchain_core.runnables import RunnableParallel, RunnablePassthrough
 from langchain_community.llms import HuggingFaceEndpoint
-import intel_extension_for_pytorch as ipex
-import torch
 
 from rag_redis.config import (
     EMBED_MODEL,
@@ -22,7 +20,6 @@ class Question(BaseModel):
 
 # Init Embeddings
 embedder = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
-embedder.client= ipex.optimize(embedder.client.eval(), dtype=torch.bfloat16)
 
 #Setup semantic cache for LLM
 from langchain.cache import RedisSemanticCache