diff --git a/crates/llm_gateway/src/stream_context.rs b/crates/llm_gateway/src/stream_context.rs index 5f50e0f7..ef00f8ec 100644 --- a/crates/llm_gateway/src/stream_context.rs +++ b/crates/llm_gateway/src/stream_context.rs @@ -159,11 +159,11 @@ impl StreamContext { // Tokenize and record token count. let token_count = tokenizer::token_count(model, json_string).unwrap_or(0); + trace!("Recorded input token count: {}", token_count); // Record the token count to metrics. self.metrics .input_sequence_length .record(token_count as u64); - trace!("Recorded input token count: {}", token_count); // Check if rate limiting needs to be applied. if let Some(selector) = self.ratelimit_selector.take() { diff --git a/crates/llm_gateway/tests/integration.rs b/crates/llm_gateway/tests/integration.rs index b18a0d29..71ec2e16 100644 --- a/crates/llm_gateway/tests/integration.rs +++ b/crates/llm_gateway/tests/integration.rs @@ -20,6 +20,11 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) { .call_proxy_on_request_headers(http_context, 0, false) .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some(":path")) .returning(Some("/v1/chat/completions")) + .expect_get_header_map_value( + Some(MapType::HttpRequestHeaders), + Some("x-arch-llm-provider"), + ) + .returning(None) .expect_get_header_map_value( Some(MapType::HttpRequestHeaders), Some("x-arch-llm-provider-hint"), @@ -36,6 +41,7 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) { Some("Authorization"), Some("Bearer secret_key"), ) + .expect_remove_header_map_value(Some(MapType::HttpRequestHeaders), Some("content-length")) .expect_get_header_map_value( Some(MapType::HttpRequestHeaders), Some("x-arch-llm-provider-hint"), @@ -48,8 +54,6 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) { .returning(Some("selector-key")) .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some("selector-key")) .returning(Some("selector-value")) - .expect_get_header_map_pairs(Some(MapType::HttpRequestHeaders)) - .returning(None) .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some(":path")) .returning(Some("/v1/chat/completions")) .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some("x-request-id")) @@ -223,8 +227,8 @@ fn llm_gateway_successful_request_to_open_ai_chat_completions() { .returning(Some(chat_completions_request_body)) .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Trace), None) - .expect_metric_record("input_sequence_length", 21) .expect_log(Some(LogLevel::Trace), None) + .expect_metric_record("input_sequence_length", 21) .expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None) .expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None) @@ -266,7 +270,7 @@ fn llm_gateway_bad_request_to_open_ai_chat_completions() { {\ \"messages\": [\ {\ - \"role\": \"system\",\ + \"role\": \"system\"\ },\ {\ \"role\": \"user\",\ @@ -283,14 +287,19 @@ fn llm_gateway_bad_request_to_open_ai_chat_completions() { ) .expect_get_buffer_bytes(Some(BufferType::HttpRequestBody)) .returning(Some(incomplete_chat_completions_request_body)) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_send_local_response( Some(StatusCode::BAD_REQUEST.as_u16().into()), None, None, None, ) - .execute_and_expect(ReturnType::Action(Action::Pause)) + .expect_log(Some(LogLevel::Trace), None) + .expect_log(Some(LogLevel::Trace), None) + .expect_metric_record("input_sequence_length", 14) + .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Debug), None) + .execute_and_expect(ReturnType::Action(Action::Continue)) .unwrap(); } diff --git a/crates/prompt_gateway/tests/integration.rs b/crates/prompt_gateway/tests/integration.rs index c5d23297..1673ac16 100644 --- a/crates/prompt_gateway/tests/integration.rs +++ b/crates/prompt_gateway/tests/integration.rs @@ -352,10 +352,10 @@ fn prompt_gateway_request_to_llm_gateway() { tool_type: ToolType::Function, function: FunctionCallDetail { name: String::from("weather_forecast"), - arguments: HashMap::from([( + arguments: Some(HashMap::from([( String::from("city"), Value::String(String::from("seattle")), - )]), + )])), }, }]), model: None, @@ -381,8 +381,8 @@ fn prompt_gateway_request_to_llm_gateway() { .expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) - .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_http_call( Some("arch_internal"), Some(vec![ @@ -400,6 +400,7 @@ fn prompt_gateway_request_to_llm_gateway() { ) .returning(Some(2)) .expect_metric_increment("active_http_calls", 1) + .expect_log(Some(LogLevel::Trace), None) .execute_and_expect(ReturnType::None) .unwrap(); @@ -499,10 +500,10 @@ fn prompt_gateway_request_no_intent_match() { tool_type: ToolType::Function, function: FunctionCallDetail { name: String::from("weather_forecast"), - arguments: HashMap::from([( + arguments: Some(HashMap::from([( String::from("city"), Value::String(String::from("seattle")), - )]), + )])), }, }]), model: None, @@ -655,10 +656,10 @@ fn prompt_gateway_request_no_intent_match_default_target() { tool_type: ToolType::Function, function: FunctionCallDetail { name: String::from("weather_forecast"), - arguments: HashMap::from([( + arguments: Some(HashMap::from([( String::from("city"), Value::String(String::from("seattle")), - )]), + )])), }, }]), model: None, diff --git a/demos/shared/test_runner/run_demo_tests.sh b/demos/shared/test_runner/run_demo_tests.sh index fcb35790..6aae9052 100644 --- a/demos/shared/test_runner/run_demo_tests.sh +++ b/demos/shared/test_runner/run_demo_tests.sh @@ -15,7 +15,7 @@ do echo "starting hurl tests" pwd ls - hurl --test hurl_tests + hurl --test hurl_tests/*.hurl echo "stopping docker containers and archgw" archgw down docker compose down -v