Skip to content

Commit

Permalink
Improve logging (#209)
Browse files Browse the repository at this point in the history 
* improve logging

* fix int tests

* better

* fix more logs

* fix more

* fix int
  • Loading branch information
@adilhafeez
adilhafeez authored Oct 22, 2024
1 parent 2f374df commit ea76d85
Show file tree
Hide file tree
Showing 6 changed files with 313 additions and 303 deletions.
1 change: 0 additions & 1 deletion crates/common/src/consts.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ pub const SYSTEM_ROLE: &str = "system";
pub const USER_ROLE: &str = "user";
pub const TOOL_ROLE: &str = "tool";
pub const ASSISTANT_ROLE: &str = "assistant";
pub const GPT_35_TURBO: &str = "gpt-3.5-turbo";
pub const ARCH_FC_REQUEST_TIMEOUT_MS: u64 = 120000; // 2 minutes
pub const MODEL_SERVER_NAME: &str = "model_server";
pub const ZEROSHOT_INTERNAL_HOST: &str = "zeroshot";
Expand Down
7 changes: 4 additions & 3 deletions crates/common/src/http.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use crate::{
stats::{Gauge, IncrementingMetric},
};
use derivative::Derivative;
use log::debug;
use log::{debug, trace};
use proxy_wasm::{traits::Context, types::Status};
use serde::Serialize;
use std::{cell::RefCell, collections::HashMap, fmt::Debug, time::Duration};
Expand Down Expand Up @@ -48,9 +48,10 @@ pub trait Client: Context {
call_args: CallArgs,
call_context: Self::CallContext,
) -> Result<u32, ClientError> {
debug!(
trace!(
"dispatching http call with args={:?} context={:?}",
call_args, call_context
call_args,
call_context
);

match self.dispatch_http_call(
Expand Down
21 changes: 6 additions & 15 deletions crates/prompt_gateway/src/context.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,24 +74,15 @@ impl Context for StreamContext {
*/

if let Some(body) = self.get_http_call_response_body(0, body_size) {
#[cfg_attr(any(), rustfmt::skip)]
match callout_context.response_handler_type {
ResponseHandlerType::GetEmbeddings => {
self.embeddings_handler(body, callout_context)
}
ResponseHandlerType::ArchGuard => self.arch_guard_handler(body, callout_context),
ResponseHandlerType::ZeroShotIntent => {
self.zero_shot_intent_detection_resp_handler(body, callout_context)
}
ResponseHandlerType::Embeddings => self.embeddings_handler(body, callout_context),
ResponseHandlerType::ZeroShotIntent => self.zero_shot_intent_detection_resp_handler(body, callout_context),
ResponseHandlerType::ArchFC => self.arch_fc_response_handler(body, callout_context),
ResponseHandlerType::HallucinationDetect => {
self.hallucination_classification_resp_handler(body, callout_context)
}
ResponseHandlerType::FunctionCall => {
self.function_call_response_handler(body, callout_context)
}
ResponseHandlerType::DefaultTarget => {
self.default_target_handler(body, callout_context)
}
ResponseHandlerType::Hallucination => self.hallucination_classification_resp_handler(body, callout_context),
ResponseHandlerType::FunctionCall => self.api_call_response_handler(body, callout_context),
ResponseHandlerType::DefaultTarget =>self.default_target_handler(body, callout_context),
}
} else {
self.send_server_error(
Expand Down
71 changes: 39 additions & 32 deletions crates/prompt_gateway/src/http_context.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use common::{
http::{CallArgs, Client},
};
use http::StatusCode;
use log::{debug, warn};
use log::{debug, trace, warn};
use proxy_wasm::{traits::HttpContext, types::Action};
use serde_json::Value;

Expand All @@ -36,7 +36,7 @@ impl HttpContext for StreamContext {
self.is_chat_completions_request =
self.get_http_request_header(":path").unwrap_or_default() == CHAT_COMPLETIONS_PATH;

debug!(
trace!(
"on_http_request_headers S[{}] req_headers={:?}",
self.context_id,
self.get_http_request_headers()
Expand All @@ -60,32 +60,37 @@ impl HttpContext for StreamContext {

self.request_body_size = body_size;

debug!(
trace!(
"on_http_request_body S[{}] body_size={}",
self.context_id, body_size
self.context_id,
body_size
);

let body_bytes = match self.get_http_request_body(0, body_size) {
Some(body_bytes) => body_bytes,
None => {
self.send_server_error(
ServerError::LogicError(format!(
"Failed to obtain body bytes even though body_size is {}",
body_size
)),
None,
);
return Action::Pause;
}
};

debug!("developer => archgw: {}", String::from_utf8_lossy(&body_bytes));

// Deserialize body into spec.
// Currently OpenAI API.
let mut deserialized_body: ChatCompletionsRequest =
match self.get_http_request_body(0, body_size) {
Some(body_bytes) => match serde_json::from_slice(&body_bytes) {
Ok(deserialized) => deserialized,
Err(e) => {
self.send_server_error(
ServerError::Deserialization(e),
Some(StatusCode::BAD_REQUEST),
);
return Action::Pause;
}
},
None => {
match serde_json::from_slice(&body_bytes) {
Ok(deserialized) => deserialized,
Err(e) => {
self.send_server_error(
ServerError::LogicError(format!(
"Failed to obtain body bytes even though body_size is {}",
body_size
)),
None,
ServerError::Deserialization(e),
Some(StatusCode::BAD_REQUEST),
);
return Action::Pause;
}
Expand Down Expand Up @@ -145,7 +150,6 @@ impl HttpContext for StreamContext {
similarity_scores: None,
upstream_cluster: None,
upstream_cluster_path: None,
tool_calls: None,
};
self.get_embeddings(callout_context);
return Action::Pause;
Expand Down Expand Up @@ -201,7 +205,6 @@ impl HttpContext for StreamContext {
similarity_scores: None,
upstream_cluster: None,
upstream_cluster_path: None,
tool_calls: None,
};

if let Err(e) = self.http_call(call_args, call_context) {
Expand All @@ -212,7 +215,7 @@ impl HttpContext for StreamContext {
}

fn on_http_response_headers(&mut self, _num_headers: usize, _end_of_stream: bool) -> Action {
debug!(
trace!(
"on_http_response_headers recv [S={}] headers={:?}",
self.context_id,
self.get_http_response_headers()
Expand All @@ -224,9 +227,11 @@ impl HttpContext for StreamContext {
}

fn on_http_response_body(&mut self, body_size: usize, end_of_stream: bool) -> Action {
debug!(
trace!(
"recv [S={}] bytes={} end_stream={}",
self.context_id, body_size, end_of_stream
self.context_id,
body_size,
end_of_stream
);

if !self.is_chat_completions_request {
Expand All @@ -248,14 +253,14 @@ impl HttpContext for StreamContext {
.expect("cant get response body");

if self.streaming_response {
debug!("streaming response");
trace!("streaming response");
} else {
debug!("non streaming response");
trace!("non streaming response");
let chat_completions_response: ChatCompletionsResponse =
match serde_json::from_slice(&body) {
Ok(de) => de,
Err(e) => {
debug!(
trace!(
"invalid response: {}, {}",
String::from_utf8_lossy(&body),
e
Expand Down Expand Up @@ -316,16 +321,18 @@ impl HttpContext for StreamContext {
serde_json::Value::String(arch_state_str),
);
let data_serialized = serde_json::to_string(&data).unwrap();
debug!("arch => user: {}", data_serialized);
debug!("archgw <= developer: {}", data_serialized);
self.set_http_response_body(0, body_size, data_serialized.as_bytes());
};
}
}
}

debug!(
trace!(
"recv [S={}] total_tokens={} end_stream={}",
self.context_id, self.response_tokens, end_of_stream
self.context_id,
self.response_tokens,
end_of_stream
);

Action::Continue
Expand Down
Loading

0 comments on commit ea76d85

Please sign in to comment.