From 7c145ebb4be68610bc3bb5377b754944307d44fd Mon Sep 17 00:00:00 2001
From: Julio Martins <89476495+juliomartinsdev@users.noreply.github.com>
Date: Thu, 3 Oct 2024 16:19:48 -0300
Subject: [PATCH] add jailbreak filter result, add ContentFilterResults on
 output (#864)

* add jailbreak filter result

* add content filter results on completion output

* add profanity content filter
---
 chat.go | 40 +++++++++++++++++++++++++++-------------
 1 file changed, 27 insertions(+), 13 deletions(-)

diff --git a/chat.go b/chat.go
index 9adf280..a7dee8e 100644
--- a/chat.go
+++ b/chat.go
@@ -41,11 +41,23 @@ type Violence struct {
 	Severity string `json:"severity,omitempty"`
 }
 
+type JailBreak struct {
+	Filtered bool `json:"filtered"`
+	Detected bool `json:"detected"`
+}
+
+type Profanity struct {
+	Filtered bool `json:"filtered"`
+	Detected bool `json:"detected"`
+}
+
 type ContentFilterResults struct {
-	Hate     Hate     `json:"hate,omitempty"`
-	SelfHarm SelfHarm `json:"self_harm,omitempty"`
-	Sexual   Sexual   `json:"sexual,omitempty"`
-	Violence Violence `json:"violence,omitempty"`
+	Hate      Hate      `json:"hate,omitempty"`
+	SelfHarm  SelfHarm  `json:"self_harm,omitempty"`
+	Sexual    Sexual    `json:"sexual,omitempty"`
+	Violence  Violence  `json:"violence,omitempty"`
+	JailBreak JailBreak `json:"jailbreak,omitempty"`
+	Profanity Profanity `json:"profanity,omitempty"`
 }
 
 type PromptAnnotation struct {
@@ -338,19 +350,21 @@ type ChatCompletionChoice struct {
 	// function_call: The model decided to call a function
 	// content_filter: Omitted content due to a flag from our content filters
 	// null: API response still in progress or incomplete
-	FinishReason FinishReason `json:"finish_reason"`
-	LogProbs     *LogProbs    `json:"logprobs,omitempty"`
+	FinishReason         FinishReason         `json:"finish_reason"`
+	LogProbs             *LogProbs            `json:"logprobs,omitempty"`
+	ContentFilterResults ContentFilterResults `json:"content_filter_results,omitempty"`
 }
 
 // ChatCompletionResponse represents a response structure for chat completion API.
 type ChatCompletionResponse struct {
-	ID                string                 `json:"id"`
-	Object            string                 `json:"object"`
-	Created           int64                  `json:"created"`
-	Model             string                 `json:"model"`
-	Choices           []ChatCompletionChoice `json:"choices"`
-	Usage             Usage                  `json:"usage"`
-	SystemFingerprint string                 `json:"system_fingerprint"`
+	ID                  string                 `json:"id"`
+	Object              string                 `json:"object"`
+	Created             int64                  `json:"created"`
+	Model               string                 `json:"model"`
+	Choices             []ChatCompletionChoice `json:"choices"`
+	Usage               Usage                  `json:"usage"`
+	SystemFingerprint   string                 `json:"system_fingerprint"`
+	PromptFilterResults []PromptFilterResult   `json:"prompt_filter_results,omitempty"`
 
 	httpHeader
 }