From a2516e31650129fa9202c378f1c3923de7b6f9f2 Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil.hafeez@gmail.com>
Date: Mon, 25 Nov 2024 15:49:48 -0800
Subject: [PATCH] use fix prompt guards

---
 arch/arch_config_schema.yaml            | 23 +++++++++++++++++++++++
 arch/build_filter_image.sh              |  1 -
 build_filter_image.sh                   |  1 +
 demos/weather_forecast/arch_config.yaml |  6 ++++++
 4 files changed, 30 insertions(+), 1 deletion(-)
 delete mode 100644 arch/build_filter_image.sh
 create mode 100644 build_filter_image.sh

diff --git a/arch/arch_config_schema.yaml b/arch/arch_config_schema.yaml
index e71e7134..6287661d 100644
--- a/arch/arch_config_schema.yaml
+++ b/arch/arch_config_schema.yaml
@@ -157,6 +157,29 @@ properties:
     enum:
       - llm
       - prompt
+  prompt_guards:
+    type: object
+    properties:
+      input_guards:
+        type: object
+        properties:
+          jailbreak:
+            type: object
+            properties:
+              on_exception:
+                type: object
+                properties:
+                  message:
+                    type: string
+                additionalProperties: false
+                required:
+                  - message
+            additionalProperties: false
+            required:
+              - on_exception
+        additionalProperties: false
+        required:
+          - jailbreak
 additionalProperties: false
 required:
   - version
diff --git a/arch/build_filter_image.sh b/arch/build_filter_image.sh
deleted file mode 100644
index 75ac81ce..00000000
--- a/arch/build_filter_image.sh
+++ /dev/null
@@ -1 +0,0 @@
-docker build  -f Dockerfile .. -t katanemo/archgw
diff --git a/build_filter_image.sh b/build_filter_image.sh
new file mode 100644
index 00000000..07093b23
--- /dev/null
+++ b/build_filter_image.sh
@@ -0,0 +1 @@
+docker build  -f arch/Dockerfile . -t katanemo/archgw
diff --git a/demos/weather_forecast/arch_config.yaml b/demos/weather_forecast/arch_config.yaml
index e827fc79..35a3f430 100644
--- a/demos/weather_forecast/arch_config.yaml
+++ b/demos/weather_forecast/arch_config.yaml
@@ -35,6 +35,12 @@ llm_providers:
 system_prompt: |
   You are a helpful assistant.
 
+prompt_guards:
+  input_guards:
+    jailbreak:
+      on_exception:
+        message: Looks like you're curious about my abilities, but I can only provide assistance for weather forecasting.
+
 prompt_targets:
   - name: weather_forecast
     description: Check weather information for a given city.