use fix prompt guards (#303)

This commit is contained in:
Adil Hafeez 2024-11-25 17:16:35 -08:00 committed by GitHub
parent 6f4a57b56d
commit 9c6fcdb771
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 212 additions and 112 deletions

View file

@ -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;