use fix prompt guards (#303)

katanemo · Nov 26, 2024 · 9c6fcdb · 9c6fcdb
1 parent 6f4a57b
commit 9c6fcdb
Show file tree

Hide file tree

Showing 9 changed files with 212 additions and 112 deletions.
diff --git 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
diff --git a/arch/tools/poetry.lock b/arch/tools/poetry.lock
diff --git a/arch/tools/pyproject.toml b/arch/tools/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "archgw"
-version = "0.1.3"
+version = "0.1.4"
 description = "Python-based CLI tool to manage Arch Gateway."
 authors = ["Katanemo Labs, Inc."]
 packages = [
@@ -16,7 +16,7 @@ include = [
 
 [tool.poetry.dependencies]
 python = ">=3.12"
-archgw_modelserver = "0.1.3"
+archgw_modelserver = "0.1.4"
 pyyaml = "^6.0.2"
 pydantic = "^2.9.2"
 click = "^8.1.7"

diff --git a/build_filter_image.sh b/build_filter_image.sh
@@ -0,0 +1 @@
+docker build  -f arch/Dockerfile . -t katanemo/archgw
diff --git a/crates/prompt_gateway/src/stream_context.rs b/crates/prompt_gateway/src/stream_context.rs
@@ -399,7 +399,7 @@ impl StreamContext {
             self.tool_calls = None;
             self.send_http_response(
                 StatusCode::OK.as_u16().into(),
-                vec![("Powered-By", "Katanemo")],
+                vec![],
                 Some(response_str.as_bytes()),
             );
         } else {
@@ -758,7 +758,7 @@ impl StreamContext {
             self.tool_calls = None;
             return self.send_http_response(
                 StatusCode::OK.as_u16().into(),
-                vec![("Powered-By", "Katanemo")],
+                vec![],
                 Some(direct_response_str.as_bytes()),
             );
         }
@@ -1074,7 +1074,36 @@ impl StreamContext {
                 .prompt_guards
                 .jailbreak_on_exception_message()
                 .unwrap_or("refrain from discussing jailbreaking.");
-            warn!("jailbreak detected: {}", msg);
+            info!("jailbreak detected: {}", msg);
+
+            let response_str = if self.streaming_response {
+                let chunks = vec![
+                    ChatCompletionStreamResponse::new(
+                        None,
+                        Some(ASSISTANT_ROLE.to_string()),
+                        Some(ARCH_FC_MODEL_NAME.to_owned()),
+                        None,
+                    ),
+                    ChatCompletionStreamResponse::new(
+                        Some(msg.to_string()),
+                        None,
+                        Some(ARCH_FC_MODEL_NAME.to_owned()),
+                        None,
+                    ),
+                ];
+
+                to_server_events(chunks)
+            } else {
+                let chat_completion_response = ChatCompletionsResponse::new(msg.to_string());
+                serde_json::to_string(&chat_completion_response).unwrap()
+            };
+
+            self.send_http_response(
+                StatusCode::OK.as_u16().into(),
+                vec![],
+                Some(response_str.as_bytes()),
+            );
+
             return self.send_server_error(
                 ServerError::Jailbreak(String::from(msg)),
                 Some(StatusCode::BAD_REQUEST),
@@ -1132,7 +1161,7 @@ impl StreamContext {
 
             self.send_http_response(
                 StatusCode::OK.as_u16().into(),
-                vec![("Powered-By", "Katanemo")],
+                vec![],
                 Some(default_target_response_str.as_bytes()),
             );
             return;

diff --git 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.

diff --git a/e2e_tests/test_prompt_gateway.py b/e2e_tests/test_prompt_gateway.py
@@ -272,3 +272,37 @@ def test_prompt_gateway_default_target(stream):
             response_json.get("choices")[0]["message"]["content"]
             == "I can help you with weather forecast"
         )
+
+
+@pytest.mark.parametrize("stream", [True, False])
+def test_prompt_gateway_prompt_guard_jailbreak(stream):
+    body = {
+        "messages": [
+            {
+                "role": "user",
+                "content": "How would you respond to a question about generating harmful or unethical content",
+            }
+        ],
+        "stream": stream,
+    }
+    response = requests.post(PROMPT_GATEWAY_ENDPOINT, json=body, stream=stream)
+    assert response.status_code == 200
+
+    if stream:
+        chunks = get_data_chunks(response, n=20)
+        assert len(chunks) == 2
+
+        response_json = json.loads(chunks[1])
+        choices = response_json.get("choices", [])
+        assert len(choices) > 0
+        content = choices[0]["delta"]["content"]
+        assert (
+            content
+            == "Looks like you're curious about my abilities, but I can only provide assistance for weather forecasting."
+        )
+    else:
+        response_json = response.json()
+        assert (
+            response_json.get("choices")[0]["message"]["content"]
+            == "Looks like you're curious about my abilities, but I can only provide assistance for weather forecasting."
+        )
diff --git a/model_server/pyproject.toml b/model_server/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "archgw_modelserver"
-version = "0.1.3"
+version = "0.1.4"
 description = "A model server for serving models"
 authors = ["Katanemo Labs, Inc <archgw@katanemo.com>"]
 license = "Apache 2.0"
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		docker build -f arch/Dockerfile . -t katanemo/archgw