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

use custom tags to store user_message

Browse source
This commit is contained in:
Adil Hafeez 2024-11-15 10:44:27 -08:00
parent d3c17c7abd
commit 1508743eeb
2 changed files with 65 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -288,6 +288,31 @@ impl HttpContext for StreamContext {
Action::Continue
}
fn on_http_response_headers(&mut self, _num_headers: usize, _end_of_stream: bool) -> Action {
debug!(
"on_http_response_headers [S={}] end_stream={}",
self.context_id, _end_of_stream
);
if let Some(user_message) = self.user_message.as_ref() {
if let Some(prompt) = user_message.content.as_ref() {
debug!("setting user-message header: {}", prompt);
self.set_http_response_header("x-user-message", Some(&prompt));
}
}
let tftt_time_ms = get_current_time()
.unwrap()
.duration_since(self.start_time.unwrap())
.unwrap()
.as_millis();
let tftt_time = tftt_time_ms.to_string();
self.set_http_response_header("x-time-to-first-token", Some(&tftt_time));
Action::Continue
}
fn on_http_response_body(&mut self, body_size: usize, end_of_stream: bool) -> Action {
debug!(
"on_http_response_body [S={}] bytes={} end_stream={}",
@ -364,6 +389,7 @@ impl HttpContext for StreamContext {
}
}
llm_span.add_attribute("model".to_string(), self.llm_provider().name.to_string());
llm_span.add_event(Event::new(
"time_to_first_token".to_string(),
self.ttft_time