apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-07-02 15:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity

fixing test cases, and making sure all references to the ChatCOmpletions* objects point to the new types

Browse source
This commit is contained in:
Salman Paracha 2025-08-11 22:42:13 -07:00
parent df32c7e278
commit 7253a0f203
15 changed files with 224 additions and 838 deletions
Download patch file Download diff file Expand all files Collapse all files

7
crates/llm_gateway/src/stream_context.rs
View file

@ -352,17 +352,16 @@ impl HttpContext for StreamContext {
};
// Set the resolved model using the trait method
deserialized_body.set_model(resolved_model);
deserialized_body.set_model(resolved_model.clone());
// Extract user message for tracing
self.user_message = deserialized_body.extract_user_message();
info!(
"on_http_request_body: provider: {}, model requested (in body): {}, model selected: {}, final model: {}",
"on_http_request_body: provider: {}, model requested (in body): {}, model selected: {}",
self.llm_provider().name,
model_requested,
model_name.unwrap_or(&"None".to_string()),
deserialized_body.model(),
);
// Use provider interface for streaming detection and setup
@ -376,7 +375,7 @@ impl HttpContext for StreamContext {
// Use provider interface for text extraction (after potential mutation)
let input_tokens_str = deserialized_body.extract_messages_text();
// enforce ratelimits on ingress
if let Err(e) = self.enforce_ratelimits(&model_requested, input_tokens_str.as_str()) {
if let Err(e) = self.enforce_ratelimits(&resolved_model, input_tokens_str.as_str()) {
self.send_server_error(
ServerError::ExceededRatelimit(e),
Some(StatusCode::TOO_MANY_REQUESTS),

45
crates/llm_gateway/tests/integration.rs
View file

@ -12,7 +12,7 @@ fn wasm_module() -> String {
wasm_file.exists(),
"Run `cargo build --release --target=wasm32-wasip1` first"
);
wasm_file.to_str().unwrap().to_string()
wasm_file.to_string_lossy().to_string()
}
fn request_headers_expectations(module: &mut Tester, http_context: i32) {
@ -267,17 +267,12 @@ fn llm_gateway_bad_request_to_open_ai_chat_completions() {
.returning(Some(incomplete_chat_completions_request_body))
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Info), Some("on_http_request_body: provider: open-ai-gpt-4, model requested (in body): gpt-1, model selected: gpt-4"))
.expect_send_local_response(
Some(StatusCode::BAD_REQUEST.as_u16().into()),
None,
None,
None,
)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), Some("getting token count model=gpt-4"))
.expect_log(Some(LogLevel::Debug), Some("Recorded input token count: 13"))
.expect_metric_record("input_sequence_length", 13)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), Some("Applying ratelimit for model: gpt-4"))
.expect_log(Some(LogLevel::Debug), Some(r#"Checking limit for provider=gpt-4, with selector=Header { key: "selector-key", value: "selector-value" }, consuming tokens=13"#))
.expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None)
.execute_and_expect(ReturnType::Action(Action::Continue))
.unwrap();
}
@ -386,11 +381,11 @@ fn llm_gateway_request_not_ratelimited() {
.returning(Some(chat_completions_request_body))
// The actual call is not important in this test, we just need to grab the token_id
.expect_log(Some(LogLevel::Info), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), Some("getting token count model=gpt-4"))
.expect_log(Some(LogLevel::Debug), Some("Recorded input token count: 29"))
.expect_metric_record("input_sequence_length", 29)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), Some("Applying ratelimit for model: gpt-4"))
.expect_log(Some(LogLevel::Debug), Some(r#"Checking limit for provider=gpt-4, with selector=Header { key: "selector-key", value: "selector-value" }, consuming tokens=29"#))
.expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None)
.execute_and_expect(ReturnType::Action(Action::Continue))
.unwrap();
@ -433,11 +428,11 @@ fn llm_gateway_override_model_name() {
// The actual call is not important in this test, we just need to grab the token_id
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Info), Some("on_http_request_body: provider: open-ai-gpt-4, model requested (in body): gpt-1, model selected: gpt-4"))
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), Some("getting token count model=gpt-4"))
.expect_log(Some(LogLevel::Debug), Some("Recorded input token count: 29"))
.expect_metric_record("input_sequence_length", 29)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), Some("Applying ratelimit for model: gpt-4"))
.expect_log(Some(LogLevel::Debug), Some(r#"Checking limit for provider=gpt-4, with selector=Header { key: "selector-key", value: "selector-value" }, consuming tokens=29"#))
.expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None)
.execute_and_expect(ReturnType::Action(Action::Continue))
.unwrap();
@ -483,8 +478,8 @@ fn llm_gateway_override_use_default_model() {
Some(LogLevel::Info),
Some("on_http_request_body: provider: open-ai-gpt-4, model requested (in body): gpt-1, model selected: gpt-4"),
)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), Some("getting token count model=gpt-4"))
.expect_log(Some(LogLevel::Debug), Some("Recorded input token count: 29"))
.expect_metric_record("input_sequence_length", 29)
.expect_log(Some(LogLevel::Debug), Some("Applying ratelimit for model: gpt-4"))
.expect_log(Some(LogLevel::Debug), Some(r#"Checking limit for provider=gpt-4, with selector=Header { key: "selector-key", value: "selector-value" }, consuming tokens=29"#))
@ -530,11 +525,11 @@ fn llm_gateway_override_use_model_name_none() {
// The actual call is not important in this test, we just need to grab the token_id
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Info), Some("on_http_request_body: provider: open-ai-gpt-4, model requested (in body): none, model selected: gpt-4"))
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), Some("getting token count model=gpt-4"))
.expect_log(Some(LogLevel::Debug), Some("Recorded input token count: 29"))
.expect_metric_record("input_sequence_length", 29)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), None)
.expect_log(Some(LogLevel::Debug), Some("Applying ratelimit for model: gpt-4"))
.expect_log(Some(LogLevel::Debug), Some(r#"Checking limit for provider=gpt-4, with selector=Header { key: "selector-key", value: "selector-value" }, consuming tokens=29"#))
.expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None)
.execute_and_expect(ReturnType::Action(Action::Continue))
.unwrap();