diff --git a/crates/prompt_gateway/src/hallucination.rs b/crates/prompt_gateway/src/hallucination.rs index 5c82c1ae..a17c6f60 100644 --- a/crates/prompt_gateway/src/hallucination.rs +++ b/crates/prompt_gateway/src/hallucination.rs @@ -104,4 +104,58 @@ mod test { println!("{:?}", messages_for_halluncination); assert_eq!(messages_for_halluncination.len(), 3); } + #[test] + fn test_hallucination_message_long() { + let test_str = r#" + [ + { + "role": "system", + "model" : "gpt-3.5-turbo", + "content": "You are a helpful assistant.\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n\n{\"type\": \"function\", \"function\": {\"name\": \"headcount\", \"description\": \"Get headcount data for a region by staffing type\", \"parameters\": {\"properties\": {\"staffing_type\": {\"type\": \"str\", \"description\": \"The staffing type like contract, fte or agency\"}, \"region\": {\"type\": \"str\", \"description\": \"the geographical region for which you want headcount data.\"}}, \"required\": [\"staffing_type\", \"region\"]}}}\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n" + }, + { "role": "user", "content": "Hello" }, + { + "role": "assistant", + "model": "gpt-3.5-turbo", + "content": "Hi there!" + }, + { "role": "user", "content": "tell me about headcount data" }, + { + "role": "assistant", + "model": "Arch-Function-1.5B", + "content": "The \"headcount\" tool provides information about the number of employees in a specific region based on the type of staffing used. It requires two parameters: \"staffing_type\" and \"region\". The \"staffing_type\" parameter specifies the type of staffing, such as contract, full-time equivalent (fte), or agency. The \"region\" parameter specifies the geographical region for which you want headcount data." + }, + { "role": "user", "content": "europe" }, + { + "role": "system", + "model": "Arch-Function-1.5B", + "content": "It seems like you are asking for headcount data for Europe. Could you please specify the staffing type?" + }, + { "role": "user", "content": "fte" }, + { + "role": "assistant", + "model": "gpt-3.5-turbo", + "content": "The headcount is 50000" + }, + { "role": "user", "content": "tell me about the weather" }, + { + "role": "assistant", + "model": "Arch-Function-1.5B", + "content" : "The weather forcast tools requires 2 parameters: city and days. Please specify" + }, + { "role": "user", "content": "Seattle" }, + { + "role": "system", + "model": "Arch-Function-1.5B", + "content": "It seems like you are asking for weather data for Seattle. Could you please specify the days?" + }, + { "role": "user", "content": "7 days" } + ] + "#; + + let messages: Vec = serde_json::from_str(test_str).unwrap(); + let messages_for_halluncination = extract_messages_for_hallucination(&messages); + println!("{:?}", messages_for_halluncination); + assert_eq!(messages_for_halluncination.len(), 3); + } }