output detection for chat completions

[resoluteCoder]

Added output detection for chat completions
Design choices

• added output warning so that the user can have the same experience with input. As well as having a way to scan the warnings if there were some detections from the output.
• created DetectionResult struct to allow for re-use of similar logic such as message_detection.
• removed the message_index var in filter_chat_messages to allow for the output detection choice_index to match the choices array in the response

Example 1

input request

 "messages": [
    {
      "content": "someemail@domain.com",
      "role": "system",
      "name": "string"
    },
    {
      "content": "someemail@domain.com",
      "role": "system",
      "name": "string"
    }
  ]

input response

{
  "id": "1a4a422085da4849b98b40876103a309",
  "object": "",
  "created": 1738705014,
  "model": "Qwen/Qwen2.5-1.5B-Instruct",
  "choices": [],
  "usage": {
    "prompt_tokens": 0,
    "total_tokens": 0,
    "completion_tokens": 0
  },
  "detections": {
    "input": [
      {
        "message_index": 1,
        "results": [
          {
            "start": 0,
            "end": 20,
            "text": "someemail@domain.com",
            "detection": "EmailAddress",
            "detection_type": "pii",
            "detector_id": "regex",
            "score": 1.0
          }
        ]
      }
    ]
  },
  "warnings": [
    {
      "type": "UNSUITABLE_INPUT",
      "message": "Unsuitable input detected. Please check the detected entities on your input and try again with the unsuitable input removed."
    }
  ]
}

Example 2

output request

"detectors": {
    "input": {
      "regex": {
        "regex": ["email"]
      }
    },
    "output": {
      "regex": {
        "regex": ["ssn"]
      }
    }
  },
  "n": 3,
  "text_gen_parameters": {
    "max_new_tokens": 25
  },
  "messages": [
    {
      "content": "can you give me an example of a social security number?",
      "role": "system",
      "name": "string"
    }
  ]

output response

{
  "id": "a2e06458b58a4a58add30f3b8e6e2beb",
  "object": "",
  "created": 1738705212,
  "model": "Qwen/Qwen2.5-1.5B-Instruct",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": "Sure! Here is an example of a Social Security Number in numerical form:\n\n123-45-6789\n\nNote that Social Security Numbers are typically 9 digits long."
      },
      "logprobs": null,
      "finish_reason": "stop"
    },
    {
      "index": 1,
      "message": {
        "role": "assistant",
        "content": "Certainly! A Social Security Number (SSN) typically consists of nine digits. An example would be:\n\n123-45-6789\n\nThis is a typical layout, and the actual digits could vary. Social Security Numbers are used to identify individuals and are unique to each person."
      },
      "logprobs": null,
      "finish_reason": "stop"
    }
  ],
  "usage": {
    "prompt_tokens": 20,
    "total_tokens": 122,
    "completion_tokens": 102
  },
  "detections": {
    "output": [
      {
        "choice_index": 0,
        "results": [
          {
            "start": 73,
            "end": 84,
            "text": "123-45-6789",
            "detection": "SocialSecurity",
            "detection_type": "pii",
            "detector_id": "regex",
            "score": 1.0
          }
        ]
      },
      {
        "choice_index": 1,
        "results": [
          {
            "start": 99,
            "end": 110,
            "text": "123-45-6789",
            "detection": "SocialSecurity",
            "detection_type": "pii",
            "detector_id": "regex",
            "score": 1.0
          }
        ]
      }
    ]
  },
  "warnings": [
    {
      "type": "UNSUITABLE_OUTPUT",
      "message": "Unsuitable output detected."
    }
  ]
}

[evaline-ju]

Thanks for the contribution! Some comments, along with a couple high level things

• Formatting is failing
• Let’s remove any unused/old functions/comments/printlns
• For example 2, the output request may not match the response? If n=3, max_new_tokens=25, I would expect 3 choices and 25 generated tokens unless some parameters were being passed to the chat completions model incorrectly?
• Have we confirmed that not providing any detectors in the request provides an expected passthrough response? (i.e. same fields are persisted)

[resoluteCoder]

@evaline-ju this is without any detectors

request

{
  "model": "Qwen/Qwen2.5-1.5B-Instruct",
  "guardrail_config": {
    "input": {
      "masks": [],
      "models": {
      }
    }
  },
  "detectors": {
    "input": {
    },
    "output": {
    }
  },
  "max_completion_tokens": 100,
  "n": 3,
  "text_gen_parameters": {
    "max_new_tokens": 25
  },
  "messages": [
    {
      "content": "hello how are you?",
      "role": "system",
      "name": "string"
    }
  ]
}

response

{
  "id": "chatcmpl-23850f08549d452da9f9a2a5b1812579",
  "object": "chat.completion",
  "created": 1738775062,
  "model": "Qwen/Qwen2.5-1.5B-Instruct",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": "Hello! I'm just a computer program, so I don't have feelings. How can I assist you today?"
      },
      "logprobs": null,
      "finish_reason": "stop"
    },
    {
      "index": 1,
      "message": {
        "role": "assistant",
        "content": "Hello! I'm not human, but a computer program. I'm here to help with questions and tasks. How can I assist you today?"
      },
      "logprobs": null,
      "finish_reason": "stop"
    },
    {
      "index": 2,
      "message": {
        "role": "assistant",
        "content": "Hello! I'm a large language model, so I don't have feelings like humans do. But I'm here to help you with anything you need! How can I assist you today?"
      },
      "logprobs": null,
      "finish_reason": "stop"
    }
  ],
  "usage": {
    "prompt_tokens": 13,
    "total_tokens": 106,
    "completion_tokens": 93
  }
}

[gkumbhat]

Thanks @resoluteCoder for the contributions.. I left a few questions and comments

[gkumbhat]

nit: since this function has gotten quite big, can we split out all of this logic (fetching the output detection part) into separate function ?

[resoluteCoder]

Sounds good. I will break out the output detections into its own function.

[gkumbhat]

nit: since this is basically same as the processing done at input time, can we move this logic into a separate function and then reuse in both input and output time?

[resoluteCoder]

I agree that they are similar and we could break that into a sort_detections function. However when I was doing just that I remembered in the input portion I returned the InputDetectionResult instead of the DetectionResult effectively casting it to the appropriate type.

If I do that with the output section I can remove where I do the conversion on line 281-287.

or of course I can remove the "casting" in those and take out the similar functionality and then loop through and "cast" them to their appropriate type.

I'm good with whichever. Regardless I did kind of mix and match those solutions 😅

Here is what it would look like

Option 1 - more iterating with concise functionality

fn sort_detections(mut detections: Vec<DetectionResult>) -> Vec<DetectionResult> {
    detections.sort_by_key(|value| value.index);

    detections
        .into_iter()
        .map(|mut detection| {
            let last_idx = detection.results.len();
            // sort detection by starting span, if span is not present then move to the end of the message
            detection.results.sort_by_key(|r| match r {
                GuardrailDetection::ContentAnalysisResponse(value) => value.start,
                _ => last_idx,
            });
            detection
        })
        .collect::<Vec<_>>()
}

    detections: Some(ChatDetections {
                        input: detections
                            .into_iter()
                            .map(|detection_result| InputDetectionResult {
                                message_index: detection_result.index,
                                results: detection_result.results,
                            })
                            .collect(),
                        output: vec![],
                    }),

Option 2 - less iterating, repetitive functionality

                let detections = input_detections
                    .into_iter()
                    .map(|mut detection| {
                        let last_idx = detection.results.len();
                        // sort detection by starting span, if span is not present then move to the end of the message
                        detection.results.sort_by_key(|r| match r {
                            GuardrailDetection::ContentAnalysisResponse(value) => value.start,
                            _ => last_idx,
                        });
                        InputDetectionResult {
                            message_index: detection.index,
                            results: detection.results,
                        }
                    })
                    .collect::<Vec<_>>();

[resoluteCoder]

@gkumbhat went ahead and choose option 1 and separated the output detections and sort detection logic in their own function

[declark1]

Thanks for the contribution @resoluteCoder. While reviewing this, I ended up having a lot of nit comments around cleanups and restructuring the processing steps, including:

• simplify function names, e.g.
  • message_detections() -> detect()
  • detector_chunk_task() -> chunk()
• chunk() should be refactored to take a list of chunker_ids and messages, returning a map of chunker_id->chunks, rather than a map of detectors as multiple detectors can use the same chunker
• detect() should only perform detection
  • preprocessing (only applicable to input detections) should be moved out of detect()
  • pass chunks to detect() instead of messages
  • return detections sorted from detect() rather than subsequently calling sort_detections()
• drop handle_output_detections() and process in match block for consistency with input detections
• mutate ChatCompletion instead of building a new object
• misc other things

I experimented with these changes and more on a fork here. Since a lot of my suggested changes are applicable to both input/output detection, I won't request the changes here. We also have a broader refactor in the works to cleanup and reuse code across handlers, which will replace some of this code with common code (e.g. the separate chunk() and detect() here won't be needed making some of these comments redundant). I'll open a separate issue/PR for items that don't overlap with the refactor. This review was still very helpful for designing common code. Thanks.

[gkumbhat]

Thanks @resoluteCoder for making the changes. Agree with some of the suggestions that @declark1 made. But if we can make those in separate PR, thats fine too.

[evaline-ju]

LGTM